Add /check-sources endpoint

This endpoint is useful for indicating whether there are sources available

api/contextqa/models/schemas.py

@@ -1,49 +1,62 @@
 # pylint: disable=E0611
 from enum import Enum
-from typing import Annotated
+from typing import Annotated, Literal
 
 from pydantic import BaseModel, Field
 
 
 class SimilarityProcessor(str, Enum):
+    """Enum representing the supported vector stores
+
+    Note that the LOCAL identifier refers to ChromaDB
+    """
+
     LOCAL = "local"
     PINECONE = "pinecone"
 
 
 class SourceFormat(str, Enum):
+    """Enum representing the supported file formats"""
+
     PDF = "pdf"
     TXT = "txt"
     CSV = "csv"
 
 
 class Source(BaseModel):
+    """Source returned as metadata in QA sessions"""
+
     title: str
     format_: Annotated[SourceFormat, Field(alias="format")]
     content: str | list
 
 
-class LLMResult(BaseModel):
-    response: str
+class SourceStatus(BaseModel):
+    """Response model returning the status of data sources"""
 
+    status: Literal["ready", "empty"]
 
-class QAResult(LLMResult):
-    sources: list[Source]
+    @classmethod
+    def from_count_status(cls, status_flag: bool) -> "SourceStatus":
+        """Get instance given the status flag"""
+        status = "ready" if status_flag else "empty"
+        return cls(status=status)
 
 
-class LLMRequestBodyBase(BaseModel):
-    separator: str = Field(description="Separator to use for the text splitting", default=".")
-    chunk_size: int = Field(description="size of each splitted chunk", default=100)
-    chunk_overlap: int = 50
+class LLMResult(BaseModel):
+    """LLM chat response object"""
+
+    response: str
 
 
 class LLMContextQueryRequest(BaseModel):
+    """QA session request object"""
+
     question: str
 
 
 class LLMQueryRequest(BaseModel):
+    """Chat request object"""
+
     message: str
     internet_access: bool = False
-
-
-class LLMQueryRequestBody(LLMRequestBodyBase):
-    query: str = Field(description="The query we want the llm to respond", min_length=10)

api/contextqa/routes/qa.py

@@ -5,11 +5,7 @@
 from sqlalchemy.orm import Session
 
 from contextqa import context, get_logger
-from contextqa.models.schemas import (
-    LLMResult,
-    SimilarityProcessor,
-    LLMContextQueryRequest,
-)
+from contextqa.models.schemas import LLMResult, SimilarityProcessor, SourceStatus, LLMContextQueryRequest
 from contextqa.routes.dependencies import get_db
 from contextqa.utils.exceptions import VectorDBConnectionError, DuplicatedSourceError
 
@@ -70,3 +66,15 @@ async def qa(params: LLMContextQueryRequest):
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail={"message": "ContextQA server did not process the request successfully", "cause": str(ex)},
         ) from ex
+
+
+@router.get("/check-sources")
+async def check_sources(session: Annotated[Session, Depends(get_db)]):
+    try:
+        status_flag = context.sources_exists(session)
+        return SourceStatus.from_count_status(status_flag)
+    except Exception as ex:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={"message": "ContextQA could not get the results from the DB", "cause": str(ex)},
+        ) from ex

api/contextqa/services/context.py

@@ -22,6 +22,7 @@
 
 from contextqa import get_logger, settings
 from contextqa.models.schemas import LLMResult, SimilarityProcessor, SourceFormat
+from contextqa.models.orm import Source
 from contextqa.utils import memory, prompts
 from contextqa.utils.exceptions import VectorDBConnectionError
 from contextqa.utils.sources import check_digest, get_not_seen_chunks
@@ -242,3 +243,18 @@ def get_setter(processor: SimilarityProcessor | None = None) -> LLMContextManage
             return LocalManager()
         case SimilarityProcessor.PINECONE:
             return PineconeManager()
+
+
+def sources_exists(session: Session) -> bool:
+    """Check if there is at least one source available
+
+    Parameters
+    ----------
+    session : Session
+        sqlalchemy session
+
+    Returns
+    -------
+    bool
+    """
+    return session.query(Source.id).limit(1).count() > 0