Dev - 1.4.0

api/contextqa/__init__.py

@@ -10,8 +10,10 @@
 def get_logger() -> logging.Logger:
     return logging.getLogger("contextqa")
 
+
 class AppSettings(BaseSettings):
     """Project settings"""
+
     default_collection: str = "contextqa-default"
     tmp_separator: str = ":::sep:::"
     media_home: Path = Path(".media/")
@@ -40,7 +42,7 @@ def validate_media_path(cls, value: Path) -> Path:
         """validator for media path"""
         value.mkdir(parents=True, exist_ok=True)
         return value
-    
+
     @property
     def sqlalchemy_url(self) -> str:
         """sqlalchemy url built either from the sqlite url or the credential of a specific mysql server"""
@@ -52,9 +54,6 @@ def sqlalchemy_url(self) -> str:
         if extras := self.mysql_extra_args:
             uri += extras
         return uri
-
-
-
 
 
 settings = AppSettings()

api/contextqa/models/schemas.py

@@ -1,49 +1,69 @@
 # 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 | dict
+    content: str | list
+
+
+class SourceStatus(BaseModel):
+    """Response model returning the status of data sources"""
+
+    status: Literal["ready", "empty"]
+
+    @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 LLMResult(BaseModel):
-    response: str
+    """LLM chat response object"""
 
+    response: str
 
-class QAResult(LLMResult):
-    sources: list[Source]
 
+class IngestionResult(BaseModel):
+    """Result of the ingestion process"""
 
-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
+    completed: int
+    skipped_files: list[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/conversational.py

@@ -1,22 +1,21 @@
 # pylint: disable=C0413
 from fastapi import APIRouter, HTTPException, status
+from fastapi.responses import StreamingResponse
 
 from contextqa import chat
-from contextqa.models.schemas import (
-    LLMResult,
-    LLMQueryRequest,
-)
+from contextqa.models.schemas import LLMQueryRequest
 
 router = APIRouter()
 
 
-@router.post("/", response_model=LLMResult)
-def get_answer(params: LLMQueryRequest):
+@router.post("/")
+async def get_answer(params: LLMQueryRequest):
     """
     Provide a message and receive a response from the LLM
     """
     try:
-        return chat.qa_service(params)
+        generator = chat.qa_service(params)
+        return StreamingResponse(generator, media_type="text/event-stream")
     except Exception as ex:
         raise HTTPException(
             status_code=status.HTTP_424_FAILED_DEPENDENCY, detail={"message": "Something went wrong", "cause": str(ex)}

api/contextqa/routes/dependencies.py

@@ -1,6 +1,7 @@
 from typing import Generator
 
 from contextqa.services.db import SessionLocal
+from sqlalchemy.orm import scoped_session
 
 
 def get_db() -> Generator:
@@ -12,7 +13,7 @@ def get_db() -> Generator:
         db session
     """
     try:
-        session = SessionLocal()
+        session = scoped_session(SessionLocal)
         yield session
         session.commit()
     except:

api/contextqa/routes/qa.py

@@ -1,15 +1,11 @@
 from typing import Annotated
 
 from fastapi import APIRouter, HTTPException, UploadFile, Depends, status
+from fastapi.responses import StreamingResponse
 from sqlalchemy.orm import Session
 
 from contextqa import context, get_logger
-from contextqa.models.schemas import (
-    LLMResult,
-    QAResult,
-    SimilarityProcessor,
-    LLMContextQueryRequest,
-)
+from contextqa.models.schemas import SimilarityProcessor, SourceStatus, LLMContextQueryRequest, IngestionResult
 from contextqa.routes.dependencies import get_db
 from contextqa.utils.exceptions import VectorDBConnectionError, DuplicatedSourceError
 
@@ -19,15 +15,16 @@
 router = APIRouter()
 
 
-@router.post("/ingest/", response_model=LLMResult)
-def ingest_source(document: UploadFile, session: Annotated[Session, Depends(get_db)]):
+@router.post("/ingest/", response_model=IngestionResult)
+def ingest_source(documents: list[UploadFile], session: Annotated[Session, Depends(get_db)]):
     """
     Ingest a data source into the vector database
     """
     try:
-        context_setter = context.get_setter(SimilarityProcessor.LOCAL)
+        context_manager = context.get_setter(SimilarityProcessor.LOCAL)
+        processor = context.BatchProcessor(manager=context_manager)
         # pylint: disable=E1102
-        return context_setter.persist(document.filename, document.file, session)
+        return processor.persist(documents, session)
     except DuplicatedSourceError as ex:
         raise HTTPException(
             status_code=status.HTTP_409_CONFLICT,
@@ -55,17 +52,29 @@ def ingest_source(document: UploadFile, session: Annotated[Session, Depends(get_
         ) from ex
 
 
-@router.post("/", response_model=QAResult)
-def qa(params: LLMContextQueryRequest):
+@router.post("/")
+async def qa(params: LLMContextQueryRequest):
     """
     Perform a QA process against the documents you have ingested
     """
     try:
         context_setter = context.get_setter()
-        # pylint: disable=E1102
-        return context_setter.load_and_respond(params.question)
+        generator = context_setter.load_and_respond(params.question)
+        return StreamingResponse(generator, media_type="text/event-stream")
     except Exception as ex:
         raise HTTPException(
             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/chat.py

@@ -1,4 +1,9 @@
+from typing import AsyncGenerator
+
 from langchain.agents import initialize_agent, AgentType, Agent
+from langchain.callbacks import AsyncIteratorCallbackHandler
+from langchain.callbacks.base import AsyncCallbackHandler
+from langchain.callbacks.streaming_aiter_final_only import AsyncFinalIteratorCallbackHandler
 from langchain.chat_models import ChatOpenAI
 from langchain.chains import ConversationChain
 from langchain.chains.conversation.prompt import DEFAULT_TEMPLATE
@@ -11,8 +16,10 @@
 
 from contextqa import settings
 from contextqa.agents.tools import searcher
-from contextqa.models.schemas import LLMResult, LLMQueryRequest
+from contextqa.models.schemas import LLMQueryRequest
 from contextqa.utils import memory, prompts
+from contextqa.agents.tools import searcher
+from contextqa.utils.streaming import stream
 
 
 _MESSAGES = [
@@ -30,7 +37,7 @@
 ]
 
 
-def get_llm_assistant(internet_access: bool) -> ConversationChain | Agent:
+def get_llm_assistant(internet_access: bool) -> tuple[ConversationChain | Agent, AsyncCallbackHandler]:
     """Return certain LLM assistant based on the system configuration
 
     Parameters
@@ -40,40 +47,44 @@ def get_llm_assistant(internet_access: bool) -> ConversationChain | Agent:
 
     Returns
     -------
-    ConversationChain | Agent
+    ConversationChain | Agent, AsyncCallbackHandler
     """
-    llm = ChatOpenAI(temperature=0)
 
     if internet_access:
-        return initialize_agent(
-            [searcher],
-            llm=llm,
-            agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
-            memory=memory.Redis("default", internet_access=True),
-            verbose=settings.debug,
-            agent_kwargs={
-                # "output_parser": CustomOP(),
-                # "format_instructions": prompts.CONTEXTQA_AGENT_TEMPLATE,
-                "prefix": prompts.PREFIX,
-            },
-            handle_parsing_errors=True,
+        callback = AsyncFinalIteratorCallbackHandler(
+            answer_prefix_tokens=["Final", "Answer", '",', "", '"', "action", "_input", '":', '"']
         )
+        llm = ChatOpenAI(temperature=0, streaming=True, callbacks=[callback])
+        return (
+            initialize_agent(
+                [searcher],
+                llm=llm,
+                agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
+                memory=memory.Redis("default", internet_access=True),
+                verbose=settings.debug,
+                agent_kwargs={"prefix": prompts.PREFIX},
+                handle_parsing_errors=True,
+            ),
+            callback,
+        )
+    callback = AsyncIteratorCallbackHandler()
+    llm = ChatOpenAI(temperature=0, streaming=True, callbacks=[callback])
     prompt = ChatPromptTemplate.from_messages(_MESSAGES)
-    return ConversationChain(llm=llm, prompt=prompt, memory=memory.Redis("default"), verbose=settings.debug)
+    return ConversationChain(llm=llm, prompt=prompt, memory=memory.Redis("default"), verbose=settings.debug), callback
 
 
-def qa_service(params: LLMQueryRequest) -> LLMResult:
+def qa_service(params: LLMQueryRequest) -> AsyncGenerator:
     """Chat with the llm
 
     Parameters
     ----------
-    params : models.LLMQueryRequest
+    params : LLMQueryRequest
         request body parameters
 
     Returns
     -------
-    models.LLMResult
-        LLM response
+    AsyncGenerator
     """
-    assistant = get_llm_assistant(params.internet_access)
-    return LLMResult(response=assistant.run(input=params.message))
+
+    assistant, callback = get_llm_assistant(params.internet_access)
+    return stream(assistant.arun(input=params.message), callback)