diff --git a/backend/app/database/weaviate/client.py b/backend/app/database/weaviate/client.py index 5ed11ede..f11268d1 100644 --- a/backend/app/database/weaviate/client.py +++ b/backend/app/database/weaviate/client.py @@ -2,10 +2,13 @@ from contextlib import asynccontextmanager from typing import AsyncGenerator import logging +import asyncio logger = logging.getLogger(__name__) _client = None +_connected = False +_client_lock = None def get_client(): @@ -15,18 +18,53 @@ def get_client(): _client = weaviate.use_async_with_local() return _client + +def _get_client_lock(): + """Get or create the client lock, binding it to the current event loop.""" + global _client_lock + if _client_lock is None: + _client_lock = asyncio.Lock() + return _client_lock + + +async def ensure_connected(): + """Ensure the client is connected. Reuses existing connection if available.""" + global _client, _connected + client = get_client() + + if not _connected or not client.is_connected(): + async with _get_client_lock(): + client = get_client() + if not _connected or not client.is_connected(): + await client.connect() + _connected = True + logger.info("Weaviate client connected") + + return client + + @asynccontextmanager async def get_weaviate_client() -> AsyncGenerator[weaviate.WeaviateClient, None]: - """Async context manager for Weaviate client.""" - client = get_client() + """Async context manager for Weaviate client with persistent connection.""" try: - await client.connect() + client = await ensure_connected() yield client except Exception as e: - logger.error(f"Weaviate client error: {str(e)}") + logger.error("Weaviate client error: %s", e) raise - finally: - try: - await client.close() - except Exception as e: - logger.warning(f"Error closing Weaviate client: {str(e)}") + + +async def close_weaviate_client(): + """Close the Weaviate client. Call this on application shutdown.""" + global _client, _connected + + async with _get_client_lock(): + if _client is not None: + try: + await _client.close() + logger.info("Weaviate client closed") + except Exception as e: + logger.warning("Error closing Weaviate client: %s", e) + finally: + _client = None + _connected = False \ No newline at end of file diff --git a/backend/main.py b/backend/main.py index b7ad80a6..38fc091b 100644 --- a/backend/main.py +++ b/backend/main.py @@ -2,6 +2,7 @@ import logging import sys from contextlib import asynccontextmanager +from app.database.weaviate.client import close_weaviate_client import uvicorn from fastapi import FastAPI, Response @@ -100,6 +101,7 @@ async def lifespan(app: FastAPI): await app_instance.start_background_tasks() yield await app_instance.stop_background_tasks() + await close_weaviate_client() api = FastAPI(title="Devr.AI API", version="1.0", lifespan=lifespan)