minor fix

app/main.py

@@ -23,29 +23,18 @@ app = FastAPI(middleware=middleware)
 
 files_dir = os.path.expanduser("~/wtp_be_files/")
 session_assistant = ChatPDF()
-# session_messages = []
-
-# async def stream_generator(agent_text_stream):
-#     for text in agent_text_stream:
-#         print(text)
-#         yield text
 
 @app.get("/query")
 async def process_input(text: str):
     if text and len(text.strip()) > 0:
         text = text.strip()
         print("PRINTING STREAM")
-        # agent_text_stream = session_assistant.ask(text)
-        # print(agent_text_stream)
-        # session_messages.append((text, True))
-        # session_messages.append((agent_text, False))
         return StreamingResponse(session_assistant.ask(text), media_type='text/event-stream')
 
 
 @app.post("/upload")
 def upload(files: list[UploadFile]):
     session_assistant.clear()
-    # session_messages = []
     
     try:
         os.makedirs(files_dir)

app/rag.py

@@ -1,25 +1,18 @@
 import os
 import logging
-
 from llama_index.core import (
     SimpleDirectoryReader,
     VectorStoreIndex,
     StorageContext,
     Settings,
     get_response_synthesizer)
-# from llama_index.core.query_engine import RetrieverQueryEngine, TransformQueryEngine
 from llama_index.core.node_parser import SentenceSplitter
 from llama_index.core.schema import TextNode, MetadataMode
-# from llama_index.core.retrievers import VectorIndexRetriever
-
-# from llama_index.core.response_synthesizers import ResponseMode
-# from transformers import AutoTokenizer
 from llama_index.core.vector_stores import VectorStoreQuery
-from llama_index.vector_stores.qdrant import QdrantVectorStore
-from qdrant_client import QdrantClient
-
 from llama_index.llms.llama_cpp import LlamaCPP
 from llama_index.embeddings.fastembed import FastEmbedEmbedding
+from llama_index.vector_stores.qdrant import QdrantVectorStore
+from qdrant_client import QdrantClient
 
 
 QDRANT_API_URL = os.getenv('QDRANT_API_URL')
@@ -31,38 +24,12 @@ logger = logging.getLogger(__name__)
 class ChatPDF:
     query_engine = None
 
-    model_url = "https://huggingface.co/Qwen/Qwen1.5-1.8B-Chat-GGUF/resolve/main/qwen1_5-1_8b-chat-q4_k_m.gguf"
-    # model_url = "https://huggingface.co/Qwen/Qwen1.5-1.8B-Chat-GGUF/resolve/main/qwen1_5-1_8b-chat-q8_0.gguf"
-    # model_url = "https://huggingface.co/microsoft/Phi-3-mini-4k-instruct-gguf/resolve/main/Phi-3-mini-4k-instruct-q4.gguf"
-
-    # def messages_to_prompt(messages):
-    #     prompt = ""
-    #     for message in messages:
-    #         if message.role == 'system':
-    #             prompt += f"<|system|>\n{message.content}</s>\n"
-    #         elif message.role == 'user':
-    #             prompt += f"<|user|>\n{message.content}</s>\n"
-    #         elif message.role == 'assistant':
-    #             prompt += f"<|assistant|>\n{message.content}</s>\n"
-
-    #     if not prompt.startswith("<|system|>\n"):
-    #         prompt = "<|system|>\n</s>\n" + prompt
-
-    #     prompt = prompt + "<|assistant|>\n"
-
-    #     return prompt
-
-    # def completion_to_prompt(completion):
-    #     return f"<|system|>\n</s>\n<|user|>\n{completion}</s>\n<|assistant|>\n"
-
-
     def __init__(self):
         self.text_parser = SentenceSplitter(chunk_size=512, chunk_overlap=24)
 
         logger.info("initializing the vector store related objects")
         # client = QdrantClient(host="localhost", port=6333)
         client = QdrantClient(url=QDRANT_API_URL, api_key=QDRANT_API_KEY)
-        # client = QdrantClient(":memory:")
         self.vector_store = QdrantVectorStore(
             client=client,
             collection_name="rag_documents",
@@ -75,7 +42,7 @@ class ChatPDF:
         )
 
         llm = LlamaCPP(
-            model_url=self.model_url,
+            model_url="https://huggingface.co/Qwen/Qwen1.5-1.8B-Chat-GGUF/resolve/main/qwen1_5-1_8b-chat-q4_k_m.gguf",
             temperature=0.1,
             max_new_tokens=256,
             context_window=3900,
@@ -132,36 +99,13 @@ class ChatPDF:
             transformations=Settings.transformations,
         )
 
-        # logger.info("configure retriever")
-        # retriever = VectorIndexRetriever(
-        #     index=index,
-        #     similarity_top_k=6,
-        #     # vector_store_query_mode="hybrid"
-        # )
-
-        # logger.info("configure response synthesizer")
-        # response_synthesizer = get_response_synthesizer(
-        #     # streaming=True,
-        #     response_mode=ResponseMode.COMPACT,
-        # )
-
-        # logger.info("assemble query engine")
-        # self.query_engine = RetrieverQueryEngine(
-        #     retriever=retriever,
-        #     response_synthesizer=response_synthesizer,
-        # )
-
         self.query_engine = index.as_query_engine(
             streaming=True,
-            # similarity_top_k=6,
+            similarity_top_k=6,
         )
 
     async def ask(self, query: str):
-        if not self.query_engine:
-            return "Please, add a PDF document first."
-
         logger.info("retrieving the response to the query")
-        # response = self.query_engine.query(str_or_query_bundle=query)
         streaming_response = self.query_engine.query(query)
         print(streaming_response)
         # return streaming_response.response_gen