feat: asgi and search service

Author: dartpain

application/api/answer/routes/search.py

@@ -1,21 +1,21 @@
 import logging
-from typing import Any, Dict, List
 
 from flask import make_response, request
 from flask_restx import fields, Resource
 
 from application.api.answer.routes.base import answer_ns
-from application.core.settings import settings
-from application.storage.db.repositories.agents import AgentsRepository
-from application.storage.db.session import db_readonly
-from application.vectorstore.vector_creator import VectorCreator
+from application.services.search_service import (
+    InvalidAPIKey,
+    SearchFailed,
+    search,
+)
 
 logger = logging.getLogger(__name__)
 
 
 @answer_ns.route("/api/search")
 class SearchResource(Resource):
-    """Fast search endpoint for retrieving relevant documents"""
+    """Fast search endpoint for retrieving relevant documents."""
 
     search_model = answer_ns.model(
         "SearchModel",
@@ -32,135 +32,24 @@ class SearchResource(Resource):
         },
     )
 
-    def _get_sources_from_api_key(self, api_key: str) -> List[str]:
-        """Get source IDs connected to the API key/agent."""
-        with db_readonly() as conn:
-            agent_data = AgentsRepository(conn).find_by_key(api_key)
-        if not agent_data:
-            return []
-
-        source_ids: List[str] = []
-        # extra_source_ids is a PG ARRAY(UUID) of source UUIDs.
-        extra = agent_data.get("extra_source_ids") or []
-        for src in extra:
-            if src:
-                source_ids.append(str(src))
-
-        if not source_ids:
-            single = agent_data.get("source_id")
-            if single:
-                source_ids.append(str(single))
-
-        return source_ids
-
-    def _search_vectorstores(
-        self, query: str, source_ids: List[str], chunks: int
-    ) -> List[Dict[str, Any]]:
-        """Search across vectorstores and return results"""
-        if not source_ids:
-            return []
-
-        results = []
-        chunks_per_source = max(1, chunks // len(source_ids))
-        seen_texts = set()
-
-        for source_id in source_ids:
-            if not source_id or not source_id.strip():
-                continue
-
-            try:
-                docsearch = VectorCreator.create_vectorstore(
-                    settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
-                )
-                docs = docsearch.search(query, k=chunks_per_source * 2)
-
-                for doc in docs:
-                    if len(results) >= chunks:
-                        break
-
-                    if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
-                        page_content = doc.page_content
-                        metadata = doc.metadata
-                    else:
-                        page_content = doc.get("text", doc.get("page_content", ""))
-                        metadata = doc.get("metadata", {})
-
-                    # Skip duplicates
-                    text_hash = hash(page_content[:200])
-                    if text_hash in seen_texts:
-                        continue
-                    seen_texts.add(text_hash)
-
-                    title = metadata.get(
-                        "title", metadata.get("post_title", "")
-                    )
-                    if not isinstance(title, str):
-                        title = str(title) if title else ""
-
-                    # Clean up title
-                    if title:
-                        title = title.split("/")[-1]
-                    else:
-                        # Use filename or first part of content as title
-                        title = metadata.get("filename", page_content[:50] + "...")
-
-                    source = metadata.get("source", source_id)
-
-                    results.append({
-                        "text": page_content,
-                        "title": title,
-                        "source": source,
-                    })
-
-                if len(results) >= chunks:
-                    break
-
-            except Exception as e:
-                logger.error(
-                    f"Error searching vectorstore {source_id}: {e}",
-                    exc_info=True,
-                )
-                continue
-
-        return results[:chunks]
-
     @answer_ns.expect(search_model)
     @answer_ns.doc(description="Search for relevant documents based on query")
     def post(self):
-        data = request.get_json()
+        data = request.get_json() or {}
 
         question = data.get("question")
         api_key = data.get("api_key")
         chunks = data.get("chunks", 5)
 
         if not question:
             return make_response({"error": "question is required"}, 400)
-
         if not api_key:
             return make_response({"error": "api_key is required"}, 400)
 
-        # Validate API key
-        with db_readonly() as conn:
-            agent = AgentsRepository(conn).find_by_key(api_key)
-        if not agent:
-            return make_response({"error": "Invalid API key"}, 401)
-
         try:
-            # Get sources connected to this API key
-            source_ids = self._get_sources_from_api_key(api_key)
-
-            if not source_ids:
-                return make_response([], 200)
-
-            # Perform search
-            results = self._search_vectorstores(question, source_ids, chunks)
-
-            return make_response(results, 200)
-
-        except Exception as e:
-            logger.error(
-                f"/api/search - error: {str(e)}",
-                extra={"error": str(e)},
-                exc_info=True,
-            )
+            return make_response(search(api_key, question, chunks), 200)
+        except InvalidAPIKey:
+            return make_response({"error": "Invalid API key"}, 401)
+        except SearchFailed:
+            logger.exception("/api/search failed")
             return make_response({"error": "Search failed"}, 500)

application/asgi.py

@@ -0,0 +1,5 @@
+from a2wsgi import WSGIMiddleware
+
+from application.app import app as flask_app
+
+asgi_app = WSGIMiddleware(flask_app)

application/requirements.txt

@@ -14,7 +14,7 @@ docx2txt==0.9
 ddgs>=8.0.0
 fast-ebook
 elevenlabs==2.43.0
-Flask==3.1.3
+Flask==3.1.1
 faiss-cpu==1.13.2
 fastmcp==3.2.4
 flask-restx==1.3.2

application/services/__init__.py

@@ -0,0 +1,150 @@
+"""Shared retrieval service used by the HTTP search route and the MCP tool.
+
+Flask-free. Raises domain exceptions (``InvalidAPIKey``, ``SearchFailed``)
+that callers translate into their own wire protocol (HTTP status codes,
+MCP error responses, etc.).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+from application.core.settings import settings
+from application.storage.db.repositories.agents import AgentsRepository
+from application.storage.db.session import db_readonly
+from application.vectorstore.vector_creator import VectorCreator
+
+logger = logging.getLogger(__name__)
+
+
+class InvalidAPIKey(Exception):
+    """The supplied ``api_key`` does not resolve to an agent."""
+
+
+class SearchFailed(Exception):
+    """Unexpected error during retrieval (e.g. DB outage). Caller maps to 5xx."""
+
+
+def _collect_source_ids(agent: Dict[str, Any]) -> List[str]:
+    """Extract the ordered list of source UUIDs to search.
+
+    Prefers ``extra_source_ids`` (PG ARRAY(UUID) of multi-source agents);
+    falls back to the legacy single ``source_id`` field.
+    """
+    source_ids: List[str] = []
+    extra = agent.get("extra_source_ids") or []
+    for src in extra:
+        if src:
+            source_ids.append(str(src))
+    if not source_ids:
+        single = agent.get("source_id")
+        if single:
+            source_ids.append(str(single))
+    return source_ids
+
+
+def _search_sources(
+    query: str, source_ids: List[str], chunks: int
+) -> List[Dict[str, Any]]:
+    """Search across each source's vectorstore and return up to ``chunks`` hits.
+
+    Per-source errors are logged and skipped so one broken index doesn't
+    take down the whole search. Results are de-duplicated by content hash.
+    """
+    if not source_ids:
+        return []
+
+    results: List[Dict[str, Any]] = []
+    chunks_per_source = max(1, chunks // len(source_ids))
+    seen_texts: set[int] = set()
+
+    for source_id in source_ids:
+        if not source_id or not source_id.strip():
+            continue
+
+        try:
+            docsearch = VectorCreator.create_vectorstore(
+                settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
+            )
+            docs = docsearch.search(query, k=chunks_per_source * 2)
+
+            for doc in docs:
+                if len(results) >= chunks:
+                    break
+
+                if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
+                    page_content = doc.page_content
+                    metadata = doc.metadata
+                else:
+                    page_content = doc.get("text", doc.get("page_content", ""))
+                    metadata = doc.get("metadata", {})
+
+                text_hash = hash(page_content[:200])
+                if text_hash in seen_texts:
+                    continue
+                seen_texts.add(text_hash)
+
+                title = metadata.get("title", metadata.get("post_title", ""))
+                if not isinstance(title, str):
+                    title = str(title) if title else ""
+
+                if title:
+                    title = title.split("/")[-1]
+                else:
+                    title = metadata.get("filename", page_content[:50] + "...")
+
+                source = metadata.get("source", source_id)
+
+                results.append(
+                    {
+                        "text": page_content,
+                        "title": title,
+                        "source": source,
+                    }
+                )
+
+            if len(results) >= chunks:
+                break
+
+        except Exception as e:
+            logger.error(
+                f"Error searching vectorstore {source_id}: {e}",
+                exc_info=True,
+            )
+            continue
+
+    return results[:chunks]
+
+
+def search(api_key: str, query: str, chunks: int = 5) -> List[Dict[str, Any]]:
+    """Resolve an agent by API key and search its sources.
+
+    Args:
+        api_key: Agent API key (the opaque string stored on
+            ``agents.key`` in Postgres).
+        query: Free-text search query.
+        chunks: Max number of hits to return.
+
+    Returns:
+        List of hit dicts with ``text``, ``title``, ``source`` keys.
+        Empty list if the agent has no sources configured.
+
+    Raises:
+        InvalidAPIKey: if ``api_key`` does not resolve to an agent.
+        SearchFailed: on unexpected DB / infrastructure errors.
+    """
+    try:
+        with db_readonly() as conn:
+            agent = AgentsRepository(conn).find_by_key(api_key)
+    except Exception as e:
+        raise SearchFailed("agent lookup failed") from e
+
+    if not agent:
+        raise InvalidAPIKey()
+
+    source_ids = _collect_source_ids(agent)
+    if not source_ids:
+        return []
+
+    return _search_sources(query, source_ids, chunks)