✨(backend) enhance Find API integration with user sub and tag
I enhance Find API integration with user access control and configuration options
This commit is contained in:
@@ -246,7 +246,7 @@ For Mistral AI models using the Etalab platform:
|
|||||||
{
|
{
|
||||||
"hrid": "mistral-medium",
|
"hrid": "mistral-medium",
|
||||||
"model_name": "mistral-medium-2508",
|
"model_name": "mistral-medium-2508",
|
||||||
"human_readable_name": "Mistral Large (Etalab)",
|
"human_readable_name": "Mistral Medium (Etalab)",
|
||||||
"provider_name": "mistral-etalab",
|
"provider_name": "mistral-etalab",
|
||||||
"profile": null,
|
"profile": null,
|
||||||
"settings": {
|
"settings": {
|
||||||
|
|||||||
@@ -166,7 +166,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
|||||||
|
|
||||||
return markdown_content
|
return markdown_content
|
||||||
|
|
||||||
def store_document(self, name: str, content: str) -> None:
|
def store_document(self, name: str, content: str, **kwargs) -> None:
|
||||||
"""
|
"""
|
||||||
Store the document content in the Albert collection.
|
Store the document content in the Albert collection.
|
||||||
This method should handle the logic to send the document content to the Albert API.
|
This method should handle the logic to send the document content to the Albert API.
|
||||||
@@ -174,6 +174,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
|||||||
Args:
|
Args:
|
||||||
name (str): The name of the document.
|
name (str): The name of the document.
|
||||||
content (str): The content of the document in Markdown format.
|
content (str): The content of the document in Markdown format.
|
||||||
|
**kwargs: Additional arguments.
|
||||||
"""
|
"""
|
||||||
response = requests.post(
|
response = requests.post(
|
||||||
urljoin(self._base_url, self._documents_endpoint),
|
urljoin(self._base_url, self._documents_endpoint),
|
||||||
@@ -188,7 +189,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
|||||||
logger.debug(response.json())
|
logger.debug(response.json())
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
|
|
||||||
async def astore_document(self, name: str, content: str) -> None:
|
async def astore_document(self, name: str, content: str, **kwargs) -> None:
|
||||||
"""
|
"""
|
||||||
Store the document content in the Albert collection.
|
Store the document content in the Albert collection.
|
||||||
This method should handle the logic to send the document content to the Albert API.
|
This method should handle the logic to send the document content to the Albert API.
|
||||||
@@ -196,6 +197,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
|||||||
Args:
|
Args:
|
||||||
name (str): The name of the document.
|
name (str): The name of the document.
|
||||||
content (str): The content of the document in Markdown format.
|
content (str): The content of the document in Markdown format.
|
||||||
|
**kwargs: Additional arguments.
|
||||||
"""
|
"""
|
||||||
async with httpx.AsyncClient(timeout=settings.ALBERT_API_TIMEOUT) as client:
|
async with httpx.AsyncClient(timeout=settings.ALBERT_API_TIMEOUT) as client:
|
||||||
response = await client.post(
|
response = await client.post(
|
||||||
|
|||||||
@@ -53,7 +53,7 @@ class BaseRagBackend(ABC):
|
|||||||
|
|
||||||
collection_ids = []
|
collection_ids = []
|
||||||
if self.collection_id:
|
if self.collection_id:
|
||||||
collection_ids.append(int(self.collection_id))
|
collection_ids.append(self.collection_id)
|
||||||
if self.read_only_collection_id:
|
if self.read_only_collection_id:
|
||||||
collection_ids.extend(
|
collection_ids.extend(
|
||||||
[int(collection_id) for collection_id in self.read_only_collection_id]
|
[int(collection_id) for collection_id in self.read_only_collection_id]
|
||||||
@@ -84,14 +84,15 @@ class BaseRagBackend(ABC):
|
|||||||
Args:
|
Args:
|
||||||
name (str): The name of the document.
|
name (str): The name of the document.
|
||||||
content_type (str): The MIME type of the document (e.g., "application/pdf").
|
content_type (str): The MIME type of the document (e.g., "application/pdf").
|
||||||
content (BytesIO): The content of the document as a BytesIO stream.
|
content (bytes): The content of the document as a bytes stream.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
str: The document content in Markdown format.
|
str: The document content in Markdown format.
|
||||||
"""
|
"""
|
||||||
raise NotImplementedError("Must be implemented in subclass.")
|
raise NotImplementedError("Must be implemented in subclass.")
|
||||||
|
|
||||||
def store_document(self, name: str, content: str) -> None:
|
@abstractmethod
|
||||||
|
def store_document(self, name: str, content: str, **kwargs) -> None:
|
||||||
"""
|
"""
|
||||||
Store the document content in the collection.
|
Store the document content in the collection.
|
||||||
This method should handle the logic to send the document content to the API.
|
This method should handle the logic to send the document content to the API.
|
||||||
@@ -99,10 +100,11 @@ class BaseRagBackend(ABC):
|
|||||||
Args:
|
Args:
|
||||||
name (str): The name of the document.
|
name (str): The name of the document.
|
||||||
content (str): The content of the document in Markdown format.
|
content (str): The content of the document in Markdown format.
|
||||||
|
**kwargs: Additional arguments. ex: "user_sub" for access control.
|
||||||
"""
|
"""
|
||||||
raise NotImplementedError("Must be implemented in subclass.")
|
raise NotImplementedError("Must be implemented in subclass.")
|
||||||
|
|
||||||
async def astore_document(self, name: str, content: str) -> None:
|
async def astore_document(self, name: str, content: str, **kwargs) -> None:
|
||||||
"""
|
"""
|
||||||
Store the document content in the collection.
|
Store the document content in the collection.
|
||||||
This method should handle the logic to send the document content to the API.
|
This method should handle the logic to send the document content to the API.
|
||||||
@@ -110,23 +112,27 @@ class BaseRagBackend(ABC):
|
|||||||
Args:
|
Args:
|
||||||
name (str): The name of the document.
|
name (str): The name of the document.
|
||||||
content (str): The content of the document in Markdown format.
|
content (str): The content of the document in Markdown format.
|
||||||
|
**kwargs: Additional arguments. ex: "user_sub" for access control.
|
||||||
"""
|
"""
|
||||||
return await sync_to_async(self.store_document)(name=name, content=content)
|
return await sync_to_async(self.store_document)(name=name, content=content, **kwargs)
|
||||||
|
|
||||||
def parse_and_store_document(self, name: str, content_type: str, content: BytesIO) -> str:
|
def parse_and_store_document(
|
||||||
|
self, name: str, content_type: str, content: bytes, **kwargs
|
||||||
|
) -> str:
|
||||||
"""
|
"""
|
||||||
Parse the document and store it in the Albert collection.
|
Parse the document and store it in the Albert collection.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
name (str): The name of the document.
|
name (str): The name of the document.
|
||||||
content_type (str): The MIME type of the document (e.g., "application/pdf").
|
content_type (str): The MIME type of the document (e.g., "application/pdf").
|
||||||
content (BytesIO): The content of the document as a BytesIO stream.
|
content (bytes): The content of the document as a bytes stream.
|
||||||
|
**kwargs: Additional arguments. ex: "user_sub" for access control.
|
||||||
"""
|
"""
|
||||||
if not self.collection_id:
|
if not self.collection_id:
|
||||||
raise RuntimeError("The RAG backend requires collection_id")
|
raise RuntimeError("The RAG backend requires collection_id")
|
||||||
|
|
||||||
document_content = self.parse_document(name, content_type, content)
|
document_content = self.parse_document(name, content_type, content)
|
||||||
self.store_document(name, document_content)
|
self.store_document(name, document_content, **kwargs)
|
||||||
return document_content
|
return document_content
|
||||||
|
|
||||||
def delete_collection(self) -> None:
|
def delete_collection(self) -> None:
|
||||||
@@ -147,22 +153,22 @@ class BaseRagBackend(ABC):
|
|||||||
def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
|
def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
|
||||||
"""
|
"""
|
||||||
Search the collection for the given query.
|
Search the collection for the given query.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
query: The search query string.
|
query: The search query string.
|
||||||
results_count: Number of results to return.
|
results_count: Number of results to return.
|
||||||
**kwargs: Additional arguments. Expected: 'session' for OIDC authentication.
|
**kwargs: Additional arguments. ex: 'session' for OIDC authentication.
|
||||||
"""
|
"""
|
||||||
raise NotImplementedError("Must be implemented in subclass.")
|
raise NotImplementedError("Must be implemented in subclass.")
|
||||||
|
|
||||||
async def asearch(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
|
async def asearch(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
|
||||||
"""
|
"""
|
||||||
Search the collection for the given query asynchronously.
|
Search the collection for the given query asynchronously.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
query: The search query string.
|
query: The search query string.
|
||||||
results_count: Number of results to return.
|
results_count: Number of results to return.
|
||||||
**kwargs: Additional arguments. Expected: 'session' for OIDC authentication.
|
**kwargs: Additional arguments. ex: 'session' for OIDC authentication.
|
||||||
"""
|
"""
|
||||||
return await sync_to_async(self.search)(query=query, results_count=results_count, **kwargs)
|
return await sync_to_async(self.search)(query=query, results_count=results_count, **kwargs)
|
||||||
|
|
||||||
|
|||||||
@@ -1,8 +1,8 @@
|
|||||||
"""Implementation of the Find API for RAG document search."""
|
"""Implementation of the Find API for RAG document search."""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
import uuid
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from multiprocessing.context import AuthenticationError
|
|
||||||
from typing import List, Optional
|
from typing import List, Optional
|
||||||
from urllib.parse import urljoin
|
from urllib.parse import urljoin
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
@@ -16,7 +16,7 @@ import requests
|
|||||||
from chat.agent_rag.constants import RAGWebResult, RAGWebResults, RAGWebUsage
|
from chat.agent_rag.constants import RAGWebResult, RAGWebResults, RAGWebUsage
|
||||||
from chat.agent_rag.document_converter.markitdown import DocumentConverter
|
from chat.agent_rag.document_converter.markitdown import DocumentConverter
|
||||||
from chat.agent_rag.document_rag_backends.base_rag_backend import BaseRagBackend
|
from chat.agent_rag.document_rag_backends.base_rag_backend import BaseRagBackend
|
||||||
from utils.oicd import refresh_access_token
|
from utils.oidc import with_fresh_access_token
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -49,14 +49,13 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri
|
|||||||
if not self.api_key:
|
if not self.api_key:
|
||||||
raise ImproperlyConfigured("FIND_API_KEY must be set in Django settings.")
|
raise ImproperlyConfigured("FIND_API_KEY must be set in Django settings.")
|
||||||
|
|
||||||
def create_collection(self, name: str, description: Optional[str] = None) -> str:
|
def create_collection(self, name: str, description: Optional[str] = None) -> uuid.UUID:
|
||||||
"""
|
"""
|
||||||
init collection_id
|
init collection_id
|
||||||
"""
|
"""
|
||||||
self.collection_id = self.collection_id or 1
|
self.collection_id = self.collection_id or str(uuid.uuid4())
|
||||||
return self.collection_id
|
return self.collection_id
|
||||||
|
|
||||||
# TODO
|
|
||||||
def delete_collection(self) -> None:
|
def delete_collection(self) -> None:
|
||||||
"""
|
"""
|
||||||
Deletion not available
|
Deletion not available
|
||||||
@@ -119,13 +118,14 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri
|
|||||||
|
|
||||||
return markdown_content
|
return markdown_content
|
||||||
|
|
||||||
def store_document(self, name: str, content: str) -> None:
|
def store_document(self, name: str, content: str, **kwargs) -> None:
|
||||||
"""
|
"""
|
||||||
index document in Find
|
index document in Find
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
name (str): The name of the document.
|
name (str): The name of the document.
|
||||||
content (str): The content of the document in Markdown format.
|
content (str): The content of the document in Markdown format.
|
||||||
|
user_sub (str): The user subject identifier for access control.
|
||||||
"""
|
"""
|
||||||
logger.debug("index document '%s' in Find", name)
|
logger.debug("index document '%s' in Find", name)
|
||||||
|
|
||||||
@@ -147,15 +147,16 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri
|
|||||||
"updated_at": timezone.now().isoformat(),
|
"updated_at": timezone.now().isoformat(),
|
||||||
"tags": [f"collection-{self.collection_id}"],
|
"tags": [f"collection-{self.collection_id}"],
|
||||||
"size": len(content.encode("utf-8")),
|
"size": len(content.encode("utf-8")),
|
||||||
"users": [], # TODO
|
"users": [user_sub],
|
||||||
"groups": [],
|
"groups": [],
|
||||||
"reach": "public",
|
"reach": "authenticated",
|
||||||
"is_active": True,
|
"is_active": True,
|
||||||
},
|
},
|
||||||
timeout=settings.ALBERT_API_TIMEOUT,
|
timeout=settings.FIND_API_TIMEOUT,
|
||||||
)
|
)
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
|
|
||||||
|
@with_fresh_access_token
|
||||||
def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
|
def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
|
||||||
"""
|
"""
|
||||||
Perform a search using the Find API.
|
Perform a search using the Find API.
|
||||||
@@ -164,28 +165,22 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri
|
|||||||
Args:
|
Args:
|
||||||
query: The search query.
|
query: The search query.
|
||||||
results_count: Number of results to return.
|
results_count: Number of results to return.
|
||||||
**kwargs: Additional arguments. Expected: 'session' containing OIDC tokens.
|
**kwargs: Additional arguments. Expected: 'session' containing OIDC tokens,
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
RAGWebResults: The search results.
|
RAGWebResults: The search results.
|
||||||
"""
|
"""
|
||||||
logger.debug("search documents in Find with query '%s'", query)
|
logger.debug("search documents in Find with query '%s'", query)
|
||||||
# TODO: factor session auth in a decorator
|
|
||||||
session = refresh_access_token(kwargs.get("session"))
|
|
||||||
oidc_access_token = session.get("oidc_access_token")
|
|
||||||
if not oidc_access_token:
|
|
||||||
raise AuthenticationError({"error": "Not authenticated"})
|
|
||||||
collection_ids = self.get_all_collection_ids()
|
|
||||||
|
|
||||||
response = requests.post(
|
response = requests.post(
|
||||||
urljoin(settings.FIND_API_URL, self.search_endpoint),
|
urljoin(settings.FIND_API_URL, self.search_endpoint),
|
||||||
headers={"Authorization": f"Bearer {oidc_access_token}"},
|
headers={"Authorization": f"Bearer {kwargs["session"].get("oidc_access_token")}"},
|
||||||
json={
|
json={
|
||||||
"q": query,
|
"q": query,
|
||||||
"tags": [f"collection:{collection_id}" for collection_id in collection_ids],
|
"tags": [f"collection-{collection_id}" for collection_id in self.get_all_collection_ids()],
|
||||||
"k": 10,
|
"k": results_count,
|
||||||
},
|
},
|
||||||
timeout=settings.ALBERT_API_TIMEOUT,
|
timeout=settings.FIND_API_TIMEOUT,
|
||||||
)
|
)
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
|
|
||||||
|
|||||||
@@ -107,7 +107,7 @@ def get_model_configuration(model_hrid: str):
|
|||||||
class AIAgentService: # pylint: disable=too-many-instance-attributes
|
class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||||
"""Service class for AI-related operations (Pydantic-AI edition)."""
|
"""Service class for AI-related operations (Pydantic-AI edition)."""
|
||||||
|
|
||||||
def __init__( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments
|
def __init__( # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||||
self,
|
self,
|
||||||
conversation: models.ChatConversation,
|
conversation: models.ChatConversation,
|
||||||
user,
|
user,
|
||||||
@@ -289,6 +289,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
|||||||
name=document.identifier,
|
name=document.identifier,
|
||||||
content_type=document.media_type,
|
content_type=document.media_type,
|
||||||
content=document_data,
|
content=document_data,
|
||||||
|
user_sub=self.user.sub,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# Remote URL
|
# Remote URL
|
||||||
@@ -300,6 +301,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
|||||||
name=document.identifier,
|
name=document.identifier,
|
||||||
content_type=document.media_type,
|
content_type=document.media_type,
|
||||||
content=document.data,
|
content=document.data,
|
||||||
|
user_sub=self.user.sub,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not document.media_type.startswith("text/"):
|
if not document.media_type.startswith("text/"):
|
||||||
|
|||||||
+25
-23
@@ -78,10 +78,10 @@ def ai_settings(request, settings):
|
|||||||
@pytest.fixture(autouse=True)
|
@pytest.fixture(autouse=True)
|
||||||
def mock_refresh_access_token():
|
def mock_refresh_access_token():
|
||||||
"""Mock refresh_access_token to bypass token refresh in tests."""
|
"""Mock refresh_access_token to bypass token refresh in tests."""
|
||||||
with mock.patch("utils.oicd.refresh_access_token") as mocked_refresh_access_token:
|
with mock.patch("utils.oidc.refresh_access_token") as mocked_refresh_access_token:
|
||||||
mock_session = Mock(spec=httpx.Client)
|
session = SessionStore()
|
||||||
mock_session.access_token = "mocked-access-token"
|
session["oidc_access_token"] = "mocked-access-token"
|
||||||
mocked_refresh_access_token.return_value = mock_session
|
mocked_refresh_access_token.return_value = session
|
||||||
yield mocked_refresh_access_token
|
yield mocked_refresh_access_token
|
||||||
|
|
||||||
|
|
||||||
@@ -103,10 +103,18 @@ def fixture_sample_pdf_content():
|
|||||||
return BytesIO(pdf_data)
|
return BytesIO(pdf_data)
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(name="mock_albert_api")
|
@pytest.fixture(name="mock_document_api")
|
||||||
def fixture_mock_albert_api():
|
def fixture_mock_document_api():
|
||||||
"""Fixture to mock the Albert API endpoints."""
|
"""Fixture to mock the Albert API endpoints."""
|
||||||
# Mock collection creation
|
# Mock collection creation
|
||||||
|
|
||||||
|
document_name = "sample.pdf"
|
||||||
|
document_content = "This is the content of the PDF."
|
||||||
|
prompt_tokens = 10
|
||||||
|
completion_tokens = 20
|
||||||
|
search_method = "semantic"
|
||||||
|
search_score = 0.9
|
||||||
|
|
||||||
responses.post(
|
responses.post(
|
||||||
"https://albert.api.etalab.gouv.fr/v1/collections",
|
"https://albert.api.etalab.gouv.fr/v1/collections",
|
||||||
json={"id": "123", "name": "test-collection"},
|
json={"id": "123", "name": "test-collection"},
|
||||||
@@ -123,7 +131,7 @@ def fixture_mock_albert_api():
|
|||||||
"metadata": {"document_name": "sample.pdf"},
|
"metadata": {"document_name": "sample.pdf"},
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20},
|
"usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens},
|
||||||
},
|
},
|
||||||
status=status.HTTP_200_OK,
|
status=status.HTTP_200_OK,
|
||||||
)
|
)
|
||||||
@@ -141,24 +149,20 @@ def fixture_mock_albert_api():
|
|||||||
json={
|
json={
|
||||||
"data": [
|
"data": [
|
||||||
{
|
{
|
||||||
"method": "semantic",
|
"method": search_method,
|
||||||
"chunk": {
|
"chunk": {
|
||||||
"id": 123,
|
"id": 123,
|
||||||
"content": "This is the content of the PDF.",
|
"content": document_content,
|
||||||
"metadata": {"document_name": "sample.pdf"},
|
"metadata": {"document_name": document_name},
|
||||||
},
|
},
|
||||||
"score": 0.9,
|
"score": search_score,
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20},
|
"usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens},
|
||||||
},
|
},
|
||||||
status=status.HTTP_200_OK,
|
status=status.HTTP_200_OK,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(name="mock_find_api")
|
|
||||||
def fixture_mock_find_api():
|
|
||||||
"""Fixture to mock the Find API endpoints."""
|
|
||||||
# Mock document indexing (Find API)
|
# Mock document indexing (Find API)
|
||||||
responses.post(
|
responses.post(
|
||||||
"https://find.api.example.com/api/v1.0/documents/index/",
|
"https://find.api.example.com/api/v1.0/documents/index/",
|
||||||
@@ -172,10 +176,10 @@ def fixture_mock_find_api():
|
|||||||
json=[
|
json=[
|
||||||
{
|
{
|
||||||
"_source": {
|
"_source": {
|
||||||
"title.fr": "sample.pdf",
|
"title.fr": document_name,
|
||||||
"content.fr": "This is the content of the PDF.",
|
"content.fr": document_content,
|
||||||
},
|
},
|
||||||
"_score": 0.9,
|
"_score": search_score,
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
status=status.HTTP_200_OK,
|
status=status.HTTP_200_OK,
|
||||||
@@ -267,8 +271,7 @@ def fixture_mock_openai_stream():
|
|||||||
def test_post_conversation_with_document_upload(
|
def test_post_conversation_with_document_upload(
|
||||||
# pylint: disable=too-many-arguments,too-many-positional-arguments
|
# pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||||
api_client,
|
api_client,
|
||||||
mock_albert_api, # pylint: disable=unused-argument
|
mock_document_api, # pylint: disable=unused-argument
|
||||||
mock_find_api, # pylint: disable=unused-argument
|
|
||||||
sample_pdf_content,
|
sample_pdf_content,
|
||||||
today_promt_date,
|
today_promt_date,
|
||||||
mock_ai_agent_service,
|
mock_ai_agent_service,
|
||||||
@@ -597,8 +600,7 @@ def test_post_conversation_with_document_upload_feature_disabled(
|
|||||||
@freeze_time()
|
@freeze_time()
|
||||||
def test_post_conversation_with_document_upload_summarize( # pylint: disable=too-many-arguments,too-many-positional-arguments # noqa: PLR0913
|
def test_post_conversation_with_document_upload_summarize( # pylint: disable=too-many-arguments,too-many-positional-arguments # noqa: PLR0913
|
||||||
api_client,
|
api_client,
|
||||||
mock_albert_api, # pylint: disable=unused-argument
|
mock_document_api, # pylint: disable=unused-argument
|
||||||
mock_find_api, # pylint: disable=unused-argument
|
|
||||||
sample_pdf_content,
|
sample_pdf_content,
|
||||||
today_promt_date,
|
today_promt_date,
|
||||||
mock_ai_agent_service,
|
mock_ai_agent_service,
|
||||||
|
|||||||
+19
-2
@@ -37,11 +37,20 @@ from chat.tests.utils import replace_uuids_with_placeholder
|
|||||||
pytestmark = pytest.mark.django_db(transaction=True)
|
pytestmark = pytest.mark.django_db(transaction=True)
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(autouse=True)
|
@pytest.fixture(
|
||||||
def ai_settings(settings):
|
autouse=True,
|
||||||
|
params=[
|
||||||
|
"chat.agent_rag.document_rag_backends.find_rag_backend.FindRagBackend",
|
||||||
|
"chat.agent_rag.document_rag_backends.albert_rag_backend.AlbertRagBackend",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def ai_settings(request, settings):
|
||||||
"""Fixture to set AI service URLs for testing."""
|
"""Fixture to set AI service URLs for testing."""
|
||||||
|
settings.RAG_DOCUMENT_SEARCH_BACKEND = request.param
|
||||||
settings.AI_BASE_URL = "https://www.external-ai-service.com/"
|
settings.AI_BASE_URL = "https://www.external-ai-service.com/"
|
||||||
settings.AI_API_KEY = "test-api-key"
|
settings.AI_API_KEY = "test-api-key"
|
||||||
|
settings.FIND_API_URL = "https://app-find"
|
||||||
|
settings.FIND_API_KEY = "find-api-key"
|
||||||
settings.AI_MODEL = "test-model"
|
settings.AI_MODEL = "test-model"
|
||||||
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful test assistant :)"
|
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful test assistant :)"
|
||||||
return settings
|
return settings
|
||||||
@@ -85,6 +94,10 @@ def test_post_conversation_with_local_pdf_document_url(
|
|||||||
json={"id": "document_id", "object": "document"},
|
json={"id": "document_id", "object": "document"},
|
||||||
status=200,
|
status=200,
|
||||||
)
|
)
|
||||||
|
responses.post(
|
||||||
|
"https://app-find/api/v1.0/documents/index/",
|
||||||
|
status=200,
|
||||||
|
)
|
||||||
|
|
||||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||||
api_client.force_authenticate(user=chat_conversation.owner)
|
api_client.force_authenticate(user=chat_conversation.owner)
|
||||||
@@ -801,6 +814,10 @@ def test_post_conversation_with_local_not_pdf_document_url(
|
|||||||
json={"id": "document_id", "object": "document"},
|
json={"id": "document_id", "object": "document"},
|
||||||
status=200,
|
status=200,
|
||||||
)
|
)
|
||||||
|
responses.post(
|
||||||
|
"https://app-find/api/v1.0/documents/index/",
|
||||||
|
status=200,
|
||||||
|
)
|
||||||
|
|
||||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||||
api_client.force_authenticate(user=chat_conversation.owner)
|
api_client.force_authenticate(user=chat_conversation.owner)
|
||||||
|
|||||||
@@ -127,7 +127,7 @@ async def _extract_and_summarize_snippets_async(query: str, url: str) -> List[st
|
|||||||
return []
|
return []
|
||||||
|
|
||||||
|
|
||||||
async def _fetch_and_store_async(url: str, document_store) -> None:
|
async def _fetch_and_store_async(url: str, document_store, **kwargs) -> None:
|
||||||
"""Fetch, extract and store text content from the URL in the document store."""
|
"""Fetch, extract and store text content from the URL in the document store."""
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -136,7 +136,7 @@ async def _fetch_and_store_async(url: str, document_store) -> None:
|
|||||||
logger.debug("Fetched document: %s", document)
|
logger.debug("Fetched document: %s", document)
|
||||||
|
|
||||||
if document:
|
if document:
|
||||||
await document_store.astore_document(url, document)
|
await document_store.astore_document(url, document, **kwargs)
|
||||||
except DocumentFetchError as e:
|
except DocumentFetchError as e:
|
||||||
logger.warning("Failed to fetch and store %s: %s", url, e)
|
logger.warning("Failed to fetch and store %s: %s", url, e)
|
||||||
# Continue with other documents
|
# Continue with other documents
|
||||||
|
|||||||
Reference in New Issue
Block a user