diff --git a/README.md b/README.md index 1a081420..64138fd5 100644 --- a/README.md +++ b/README.md @@ -100,7 +100,7 @@ The following environment variables are required to run the application: - azure: "text-embedding-3-small" (will be used as your Azure Deployment) - huggingface: "sentence-transformers/all-MiniLM-L6-v2" - huggingfacetei: "http://huggingfacetei:3000". Hugging Face TEI uses model defined on TEI service launch. - - vertexai: "text-embedding-004" + - vertexai: "gemini-embedding-001" - ollama: "nomic-embed-text" - bedrock: "amazon.titan-embed-text-v1" - google_genai: "gemini-embedding-001" @@ -119,7 +119,9 @@ The following environment variables are required to run the application: - `AWS_SECRET_ACCESS_KEY`: (Optional) needed for bedrock embeddings - `GOOGLE_API_KEY`, `GOOGLE_KEY`, `RAG_GOOGLE_API_KEY`: (Optional) Google API key for Google GenAI embeddings. Priority order: RAG_GOOGLE_API_KEY > GOOGLE_KEY > GOOGLE_API_KEY - `AWS_SESSION_TOKEN`: (Optional) may be needed for bedrock embeddings -- `GOOGLE_APPLICATION_CREDENTIALS`: (Optional) needed for Google VertexAI embeddings. This should be a path to a service account credential file in JSON format, as accepted by [langchain](https://python.langchain.com/api_reference/google_vertexai/index.html) +- `GOOGLE_APPLICATION_CREDENTIALS`: (Optional) needed for Google VertexAI embeddings. This should be a path to a service account credential file in JSON format. +- `GOOGLE_CLOUD_PROJECT`: (Optional) Google Cloud project ID, needed for VertexAI embeddings. +- `GOOGLE_CLOUD_LOCATION`: (Optional) Google Cloud region for VertexAI embeddings. Defaults to `us-central1`. - `RAG_CHECK_EMBEDDING_CTX_LENGTH` (Optional) Default is true, disabling this will send raw input to the embedder, use this for custom embedding models. Make sure to set these environment variables before running the application. You can set them in a `.env` file or as system environment variables. diff --git a/app/config.py b/app/config.py index b5f7847c..8031763a 100644 --- a/app/config.py +++ b/app/config.py @@ -257,12 +257,18 @@ def init_embeddings(provider, model): return GoogleGenerativeAIEmbeddings( model=model, - google_api_key=RAG_GOOGLE_API_KEY, + google_api_key=RAG_GOOGLE_API_KEY or None, ) elif provider == EmbeddingsProvider.GOOGLE_VERTEXAI: - from langchain_google_vertexai import VertexAIEmbeddings + from langchain_google_genai import GoogleGenerativeAIEmbeddings - return VertexAIEmbeddings(model=model) + return GoogleGenerativeAIEmbeddings( + model=model, + google_api_key=RAG_GOOGLE_API_KEY or None, + vertexai=True, + project=get_env_variable("GOOGLE_CLOUD_PROJECT", None), + location=get_env_variable("GOOGLE_CLOUD_LOCATION", "us-central1"), + ) elif provider == EmbeddingsProvider.BEDROCK: from langchain_aws import BedrockEmbeddings @@ -306,7 +312,7 @@ def init_embeddings(provider, model): "EMBEDDINGS_MODEL", "http://huggingfacetei:3000" ) elif EMBEDDINGS_PROVIDER == EmbeddingsProvider.GOOGLE_VERTEXAI: - EMBEDDINGS_MODEL = get_env_variable("EMBEDDINGS_MODEL", "text-embedding-004") + EMBEDDINGS_MODEL = get_env_variable("EMBEDDINGS_MODEL", "gemini-embedding-001") elif EMBEDDINGS_PROVIDER == EmbeddingsProvider.OLLAMA: EMBEDDINGS_MODEL = get_env_variable("EMBEDDINGS_MODEL", "nomic-embed-text") elif EMBEDDINGS_PROVIDER == EmbeddingsProvider.GOOGLE_GENAI: diff --git a/app/routes/document_routes.py b/app/routes/document_routes.py index 2fe56444..f96b67ec 100644 --- a/app/routes/document_routes.py +++ b/app/routes/document_routes.py @@ -21,7 +21,6 @@ status, ) from langchain_core.documents import Document -from langchain_core.runnables import run_in_executor from langchain_text_splitters import RecursiveCharacterTextSplitter from functools import lru_cache import asyncio @@ -145,7 +144,8 @@ async def load_file_content( loader = None try: loader, known_type, file_ext = get_loader(filename, content_type, file_path) - data = await run_in_executor(executor, lambda: list(loader.lazy_load())) + loop = asyncio.get_running_loop() + data = await loop.run_in_executor(executor, lambda: list(loader.lazy_load())) return data, known_type, file_ext finally: # Clean up temporary UTF-8 file if it was created for encoding conversion @@ -751,7 +751,8 @@ async def embed_local_file( loader, known_type, file_ext = get_loader( document.filename, document.file_content_type, file_path ) - data = await run_in_executor( + loop = asyncio.get_running_loop() + data = await loop.run_in_executor( request.app.state.thread_pool, lambda: list(loader.lazy_load()) ) diff --git a/app/services/vector_store/async_pg_vector.py b/app/services/vector_store/async_pg_vector.py index dd1371d8..f4d8381a 100644 --- a/app/services/vector_store/async_pg_vector.py +++ b/app/services/vector_store/async_pg_vector.py @@ -1,75 +1,100 @@ -from typing import Optional, List, Tuple, Dict, Any +from typing import Callable, Optional, List, Tuple, Dict, Any, TypeVar import asyncio +from concurrent.futures import Executor from langchain_core.documents import Document -from langchain_core.runnables.config import run_in_executor from .extended_pg_vector import ExtendedPgVector +T = TypeVar("T") + + class AsyncPgVector(ExtendedPgVector): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._thread_pool = None - + def _get_thread_pool(self): if self._thread_pool is None: try: - # Try to get the thread pool from FastAPI app state - import contextvars - from fastapi import Request - # This is a fallback - in practice, we'll pass the executor explicitly loop = asyncio.get_running_loop() - self._thread_pool = getattr(loop, '_default_executor', None) - except: + self._thread_pool = getattr(loop, "_default_executor", None) + except Exception: pass return self._thread_pool - + + @staticmethod + async def _run_in_executor( + executor: Executor | None, + func: Callable[..., T], + *args: Any, + **kwargs: Any, + ) -> T: + """Run a sync callable in a thread pool executor. + + Wraps the call to convert StopIteration into RuntimeError. + StopIteration cannot be set on an asyncio.Future — it raises + TypeError and leaves the Future pending forever. + """ + + def wrapper() -> T: + try: + return func(*args, **kwargs) + except StopIteration as exc: + raise RuntimeError from exc + + loop = asyncio.get_running_loop() + return await loop.run_in_executor(executor, wrapper) + async def get_all_ids(self, executor=None) -> list[str]: executor = executor or self._get_thread_pool() - return await run_in_executor(executor, super().get_all_ids) - + return await self._run_in_executor(executor, super().get_all_ids) + async def get_filtered_ids(self, ids: list[str], executor=None) -> list[str]: executor = executor or self._get_thread_pool() - return await run_in_executor(executor, super().get_filtered_ids, ids) + return await self._run_in_executor(executor, super().get_filtered_ids, ids) - async def get_documents_by_ids(self, ids: list[str], executor=None) -> list[Document]: + async def get_documents_by_ids( + self, ids: list[str], executor=None + ) -> list[Document]: executor = executor or self._get_thread_pool() - return await run_in_executor(executor, super().get_documents_by_ids, ids) + return await self._run_in_executor(executor, super().get_documents_by_ids, ids) async def delete( - self, ids: Optional[list[str]] = None, collection_only: bool = False, executor=None + self, + ids: Optional[list[str]] = None, + collection_only: bool = False, + executor=None, ) -> None: executor = executor or self._get_thread_pool() - await run_in_executor(executor, self._delete_multiple, ids, collection_only) - + await self._run_in_executor( + executor, self._delete_multiple, ids, collection_only + ) + async def asimilarity_search_with_score_by_vector( - self, - embedding: List[float], - k: int = 4, + self, + embedding: List[float], + k: int = 4, filter: Optional[Dict[str, Any]] = None, - executor=None + executor=None, ) -> List[Tuple[Document, float]]: """Async version of similarity_search_with_score_by_vector""" executor = executor or self._get_thread_pool() - return await run_in_executor( - executor, - super().similarity_search_with_score_by_vector, - embedding, - k, - filter + return await self._run_in_executor( + executor, + super().similarity_search_with_score_by_vector, + embedding, + k, + filter, ) - + async def aadd_documents( - self, - documents: List[Document], + self, + documents: List[Document], ids: Optional[List[str]] = None, executor=None, - **kwargs + **kwargs, ) -> List[str]: """Async version of add_documents""" executor = executor or self._get_thread_pool() - return await run_in_executor( - executor, - super().add_documents, - documents, - ids=ids, - **kwargs - ) \ No newline at end of file + return await self._run_in_executor( + executor, super().add_documents, documents, ids=ids, **kwargs + ) diff --git a/requirements.lite.txt b/requirements.lite.txt index 6a52ff1a..2f963cb2 100644 --- a/requirements.lite.txt +++ b/requirements.lite.txt @@ -1,9 +1,8 @@ -langchain==0.3.26 -langchain-community==0.3.27 -langchain-openai==0.3.27 -langchain-core==0.3.81 -langchain-google-genai==2.1.10 -langchain-google-vertexai==2.0.27 +langchain==1.2.10 +langchain-community==0.4.1 +langchain-openai==1.1.10 +langchain-core==1.2.16 +langchain-google-genai==4.2.0 sqlalchemy==2.0.41 python-dotenv==1.1.1 fastapi==0.115.12 @@ -24,14 +23,14 @@ python-multipart==0.0.22 aiofiles==24.1.0 rapidocr-onnxruntime==1.4.4 opencv-python-headless==4.9.0.80 -pymongo==4.6.3 -langchain-mongodb==0.2.0 +pymongo>=4.12.0,<5 +langchain-mongodb==0.11.0 cryptography==46.0.5 python-magic==0.4.27 python-pptx==1.0.2 xlrd==2.0.2 -langchain-aws==0.2.1 -boto3==1.34.144 +langchain-aws==1.3.1 +boto3>=1.42.42,<2 chardet==5.2.0 -langchain-ollama==0.3.3 +langchain-ollama==1.0.1 tenacity>=9.0.0 diff --git a/requirements.txt b/requirements.txt index 570e0c48..d5411536 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,10 @@ -langchain==0.3.26 -langchain-community==0.3.27 -langchain-openai==0.3.27 -langchain-core==0.3.81 -langchain-aws==0.2.1 -langchain-google-vertexai==2.0.27 -langchain_text_splitters==0.3.9 -boto3==1.34.144 +langchain==1.2.10 +langchain-community==0.4.1 +langchain-openai==1.1.10 +langchain-core==1.2.16 +langchain-aws==1.3.1 +langchain-text-splitters==1.1.1 +boto3>=1.42.42,<2 sqlalchemy==2.0.41 python-dotenv==1.1.1 fastapi==0.115.12 @@ -27,15 +26,15 @@ sentence_transformers==3.1.1 aiofiles==24.1.0 rapidocr-onnxruntime==1.4.4 opencv-python-headless==4.9.0.80 -pymongo==4.6.3 -langchain-mongodb==0.2.0 -langchain-ollama==0.3.3 -langchain-huggingface==0.1.0 -langchain-google-genai==2.1.10 +pymongo>=4.12.0,<5 +langchain-mongodb==0.11.0 +langchain-ollama==1.0.1 +langchain-huggingface==1.2.0 +langchain-google-genai==4.2.0 cryptography==46.0.5 python-magic==0.4.27 python-pptx==1.0.2 xlrd==2.0.2 -pydantic==2.9.2 +pydantic>=2.10.6,<3 chardet==5.2.0 tenacity>=9.0.0 diff --git a/tests/services/test_async_pg_vector.py b/tests/services/test_async_pg_vector.py new file mode 100644 index 00000000..cb280591 --- /dev/null +++ b/tests/services/test_async_pg_vector.py @@ -0,0 +1,95 @@ +import asyncio +from unittest.mock import patch, MagicMock +import pytest +from langchain_core.documents import Document +from app.services.vector_store.async_pg_vector import AsyncPgVector +from app.services.vector_store.extended_pg_vector import ExtendedPgVector + + +class DummyAsyncPgVector(AsyncPgVector): + """Subclass that skips DB initialization.""" + + def __init__(self): + # Bypass ExtendedPgVector/PGVector __init__ entirely + self._thread_pool = None + self._bind = None # Prevent AttributeError in PGVector.__del__ + + +@pytest.fixture +def store(): + return DummyAsyncPgVector() + + +@pytest.mark.asyncio +async def test_get_all_ids_dispatches_to_super(store): + with patch.object( + ExtendedPgVector, "get_all_ids", return_value=["id1", "id2"] + ) as mock: + result = await store.get_all_ids() + mock.assert_called_once_with() + assert result == ["id1", "id2"] + + +@pytest.mark.asyncio +async def test_get_filtered_ids_passes_ids(store): + with patch.object( + ExtendedPgVector, "get_filtered_ids", return_value=["id1"] + ) as mock: + result = await store.get_filtered_ids(["id1", "id2"]) + mock.assert_called_once_with(["id1", "id2"]) + assert result == ["id1"] + + +@pytest.mark.asyncio +async def test_get_documents_by_ids_passes_ids(store): + docs = [Document(page_content="test", metadata={"file_id": "id1"})] + with patch.object( + ExtendedPgVector, "get_documents_by_ids", return_value=docs + ) as mock: + result = await store.get_documents_by_ids(["id1"]) + mock.assert_called_once_with(["id1"]) + assert result == docs + + +@pytest.mark.asyncio +async def test_delete_passes_args(store): + with patch.object(ExtendedPgVector, "_delete_multiple") as mock: + await store.delete(ids=["id1"], collection_only=True) + mock.assert_called_once_with(["id1"], True) + + +@pytest.mark.asyncio +async def test_asimilarity_search_passes_args(store): + expected = [(Document(page_content="test", metadata={}), 0.9)] + with patch.object( + ExtendedPgVector, + "similarity_search_with_score_by_vector", + return_value=expected, + ) as mock: + embedding = [0.1, 0.2, 0.3] + result = await store.asimilarity_search_with_score_by_vector( + embedding, k=5, filter={"file_id": {"$eq": "id1"}} + ) + mock.assert_called_once_with(embedding, 5, {"file_id": {"$eq": "id1"}}) + assert result == expected + + +@pytest.mark.asyncio +async def test_aadd_documents_passes_args(store): + docs = [Document(page_content="test", metadata={})] + with patch.object(ExtendedPgVector, "add_documents", return_value=["id1"]) as mock: + result = await store.aadd_documents(docs, ids=["id1"]) + mock.assert_called_once_with(docs, ids=["id1"]) + assert result == ["id1"] + + +@pytest.mark.asyncio +async def test_run_in_executor_converts_stop_iteration(store): + """StopIteration can't be set on an asyncio.Future — verify it becomes RuntimeError.""" + + def raises_stop(): + raise StopIteration("exhausted") + + with patch.object(ExtendedPgVector, "get_all_ids", side_effect=raises_stop): + with pytest.raises(RuntimeError): + await store.get_all_ids()