diff --git a/backend/Dockerfile b/backend/Dockerfile index 28178ad2..d655a5ff 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,6 +1,6 @@ FROM --platform=linux/amd64 python:3.10 -RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 poppler-utils -y +RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 poppler-utils qpdf -y # Set environment variables ENV PATH=/virtualenvs/venv/bin:$PATH diff --git a/backend/modules/parsers/__init__.py b/backend/modules/parsers/__init__.py index a008523a..3e85eb91 100644 --- a/backend/modules/parsers/__init__.py +++ b/backend/modules/parsers/__init__.py @@ -3,8 +3,10 @@ from backend.modules.parsers.pdfparser_fast import PdfParserUsingPyMuPDF from backend.modules.parsers.tablepdfparser import PdfTableParser from backend.modules.parsers.textparser import TextParser +from backend.modules.parsers.tfparser import TfParser register_parser("MarkdownParser", MarkdownParser) register_parser("PdfParserFast", PdfParserUsingPyMuPDF) register_parser("TextParser", TextParser) register_parser("PdfTableParser", PdfTableParser) +register_parser("TfParser", TfParser) diff --git a/backend/modules/parsers/tablepdfparser.py b/backend/modules/parsers/tablepdfparser.py index d689fbca..4d62b86e 100644 --- a/backend/modules/parsers/tablepdfparser.py +++ b/backend/modules/parsers/tablepdfparser.py @@ -82,6 +82,18 @@ async def get_chunks(self, filepath, metadata, *args, **kwargs): ) ] table_docs.extend(tab_doc) + if table_data: + tab_doc = [ + Document( + page_content=table_data, + metadata={ + "page_num": page.page_number, + "type": "table", + "table_num": ix, + }, + ) + ] + table_docs.extend(tab_doc) text = page.text diff --git a/backend/modules/parsers/tfparser.py b/backend/modules/parsers/tfparser.py new file mode 100644 index 00000000..069ee5cd --- /dev/null +++ b/backend/modules/parsers/tfparser.py @@ -0,0 +1,158 @@ +import os +from typing import Optional + +import fitz +import requests +from langchain.docstore.document import Document + +from backend.logger import logger +from backend.modules.parsers.parser import BaseParser +from backend.settings import settings + + +class TfParser(BaseParser): + """ + TfParser is a multi-modal parser class for deep extraction of pdf documents. + Requires a running instance of the TfParser service that has access to TFLLM Gateway + """ + + supported_file_extensions = [".pdf"] + + def __init__(self, max_chunk_size: int = 1000, *args, **kwargs): + """ + Initializes the TfParser object. + """ + self.max_chunk_size = max_chunk_size + self.tf_service_url = settings.TF_PARSER + + async def _send_file_request(self, payload: dict, endpoint: str): + """ + Sends a POST request to the TfParser service. + """ + response = requests.post( + self.tf_service_url.rstrip("/") + endpoint, + files=payload, + ) + if "error" in response: + print(f"Error: {response.json()['error']}") + return None + return response + + async def _send_text_request(self, payload: dict, endpoint: str): + """ + Sends a POST request to the TfParser service. + """ + response = requests.post( + self.tf_service_url.rstrip("/") + endpoint, json=payload + ) + if "error" in response: + print(f"Error: {response.json()['error']}") + return None + return response + + async def get_chunks( + self, filepath: str, metadata: Optional[dict] = None, *args, **kwargs + ): + """ + Asynchronously extracts text from a PDF file and returns it in chunks. + """ + if not filepath.endswith(".pdf"): + print("Invalid file extension. TfParser only supports PDF files.") + return [] + page_texts = list() + final_texts = list() + + try: + # Open the PDF file using pdfplumber + doc = fitz.open(filepath) + + # get file path & name + head, tail = os.path.split(filepath) + + for page in doc: + try: + page_number = page.number + 1 + print(f"\n\nProcessing page {page_number}...") + + content = fitz.open() + # copy over current page + content.insert_pdf(doc, from_page=page.number, to_page=page.number) + # save the page to a temporary file + temp_file = os.path.join(head, f"{tail}-{page_number}.pdf") + content.save(temp_file) + content.close() + + # send the page to the TfParser service + with open(temp_file, "rb") as f: + response = await self._send_file_request( + payload={ + "file": f, + }, + endpoint="/tf-parse-pdf", + ) + # Parse the response + response = response.json() + + if "error" not in response: + for res in response: + page_content = res.get("page_content").strip() + if ( + page_content != "" + and page_content is not None + and page_content != " " + ): + metadata = res.get("metadata", {}) + metadata["page_number"] = page_number + metadata["source"] = tail + final_texts.append( + Document( + page_content=page_content, + metadata=metadata, + ) + ) + page_texts.append(page_content) + print( + f"Page Content: {page_content}, \nmetadata-pg-no: {metadata['page_number']}, metadata-type: {metadata['type']}" + ) + else: + print(f"Error in Page: {response['error']}") + + # remove the temporary file + print("Removing temp file...") + os.remove(temp_file) + except Exception as e: + print(f"Exception in Page: {e}") + # remove the temporary file + print("Removing temp file...") + os.remove(temp_file) + continue + + document_text = " ".join(page_texts) + if document_text: + print("\n\nProcessing combined doc...") + response = await self._send_text_request( + payload={"text": document_text, "file_name": tail}, + endpoint="/tf-get-response", + ) + response = response.json() + if "error" not in response: + page_content = response.get("page_content").strip() + metadata = response.get("metadata", {}) + if ( + page_content != "" + and page_content is not None + and page_content != " " + ): + final_texts.append( + Document( + page_content=page_content, + metadata=metadata, + ) + ) + print(f"Page Content: {page_content}") + else: + print(f"Error: {response['error']}") + return final_texts + except Exception as e: + print(f"Ultimate Exception: {e}") + return final_texts diff --git a/backend/modules/query_controllers/__init__.py b/backend/modules/query_controllers/__init__.py index f41c2c80..a504e61a 100644 --- a/backend/modules/query_controllers/__init__.py +++ b/backend/modules/query_controllers/__init__.py @@ -1,4 +1,6 @@ from backend.modules.query_controllers.example.controller import ExampleQueryController from backend.modules.query_controllers.query_controller import register_query_controller +from backend.modules.query_controllers.summary.controller import SummaryQueryController register_query_controller("default", ExampleQueryController) +register_query_controller("summary", SummaryQueryController) diff --git a/backend/modules/query_controllers/example/controller.py b/backend/modules/query_controllers/example/controller.py index 60d857e5..b24a50b0 100644 --- a/backend/modules/query_controllers/example/controller.py +++ b/backend/modules/query_controllers/example/controller.py @@ -19,6 +19,7 @@ from backend.modules.query_controllers.example.payload import ( QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD, QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, + QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_PAYLOAD, QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_MMR_PAYLOAD, QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE_PAYLOAD, @@ -40,24 +41,25 @@ EXAMPLES = { "vector-store-similarity": QUERY_WITH_VECTOR_STORE_RETRIEVER_PAYLOAD, - "multi-query-similarity-threshold": QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, + "vector-store-similarity-threshold": QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, + "contexual-compression-multi-query-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, + "contextual-compression-multi-query-similarity-threshold": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, # Keeping these for future use: + # "contexual-compression-similarity-threshold": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE_PAYLOAD, # "vector-store-mmr": QUERY_WITH_VECTOR_STORE_RETRIEVER_MMR_PAYLOAD, - # "vector-store-similarity-threshold": QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, # "contexual-compression-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_PAYLOAD, - # "contexual-compression-similarity-threshold": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE_PAYLOAD, - # "multi-query-similarity": QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, # "multi-query-mmr": QUERY_WITH_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD, - # "contexual-compression-multi-query-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, } -if settings.LOCAL: - EXAMPLES.update( - { - "contexual-compression-mmr": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_MMR_PAYLOAD, - "contexual-compression-multi-query-mmr": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD, - } - ) +# if settings.LOCAL: +# EXAMPLES.update( +# { +# "contexual-compression-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_PAYLOAD, +# "contexual-compression-multi-query-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, +# "multi-query-similarity": QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, +# "multi-query-similarity-threshold": QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, +# } +# ) @query_controller("/example-app") diff --git a/backend/modules/query_controllers/example/payload.py b/backend/modules/query_controllers/example/payload.py index 0c28e21c..f936b54b 100644 --- a/backend/modules/query_controllers/example/payload.py +++ b/backend/modules/query_controllers/example/payload.py @@ -13,7 +13,7 @@ } QUERY_WITH_VECTOR_STORE_RETRIEVER_PAYLOAD = { - "summary": "vector store retriever with similarity search", + "summary": "search with similarity", "description": """ Requires k in search_kwargs for similarity search. search_type can either be similarity or mmr or similarity_score_threshold.""", @@ -42,7 +42,7 @@ } QUERY_WITH_VECTOR_STORE_RETRIEVER_MMR_PAYLOAD = { - "summary": "vector store retriever with mmr", + "summary": "search with mmr", "description": """ Requires k and fetch_k in search_kwargs for mmr support depends on vector db. search_type can either be similarity or mmr or similarity_score_threshold""", @@ -68,7 +68,7 @@ } QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY_SCORE_PAYLOAD = { - "summary": "vector store retriever with threshold score", + "summary": "search with threshold score", "description": """ Requires score_threshold float (0~1) in search kwargs. search_type can either be similarity or mmr or similarity_score_threshold.""", @@ -89,15 +89,15 @@ "retriever_config": { "compressor_model_provider": "mixbread-ai", "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", - "top_k": 5, + "top_k": 7, "search_type": "similarity", - "search_kwargs": {"k": 20}, + "search_kwargs": {"k": 25}, }, "stream": False, } QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_PAYLOAD = { - "summary": "contextual compression (re-ranker) retriever with similariy search", + "summary": "similariy search + re-ranking", "description": """ Requires k in search_kwargs for similarity search. search_type can either be similarity or mmr or similarity_score_threshold.""", @@ -119,10 +119,10 @@ "retriever_config": { "compressor_model_provider": "mixbread-ai", "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", - "top_k": 5, + "top_k": 7, "search_type": "mmr", "search_kwargs": { - "k": 20, + "k": 25, "fetch_k": 30, }, }, @@ -130,7 +130,7 @@ } QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_MMR_PAYLOAD = { - "summary": "contextual compression (re-ranker) retriever with mmr", + "summary": "mmr + re-ranking", "description": """ Requires k and fetch_k in search kwargs for mmr. search_type can either be similarity or mmr or similarity_score_threshold. @@ -154,7 +154,7 @@ "retriever_config": { "compressor_model_provider": "mixbread-ai", "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", - "top_k": 5, + "top_k": 7, "search_type": "similarity_score_threshold", "search_kwargs": {"score_threshold": 0.7}, }, @@ -162,7 +162,7 @@ } QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE_PAYLOAD = { - "summary": "contextual compression (re-ranker) retriever with threshold score", + "summary": "threshold score + re-ranking", "description": """ Requires score_threshold float (0~1) in search kwargs for similarity search. search_type can either be similarity or mmr or similarity_score_threshold. @@ -195,7 +195,7 @@ } QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD = { - "summary": "multi-query retriever with similarity search", + "summary": "multi-query + similarity search", "description": """ Typically used for complex user queries. search_type can either be similarity or mmr or similarity_score_threshold.""", @@ -230,7 +230,7 @@ } QUERY_WITH_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD = { - "summary": "multi-query retriever with mmr", + "summary": "multi-query + mmr", "description": """ Requires k and fetch_k in search_kwargs for mmr. search_type can either be similarity or mmr or similarity_score_threshold.""", @@ -261,7 +261,7 @@ } QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD = { - "summary": "multi-query retriever with similarity score threshold", + "summary": "multi-query + threshold score", "description": """ Typically used for complex user queries. Requires score_threshold float (0~1) in search kwargs. @@ -284,10 +284,10 @@ "retriever_config": { "compressor_model_provider": "mixbread-ai", "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", - "top_k": 5, + "top_k": 7, "search_type": "mmr", "search_kwargs": { - "k": 20, + "k": 25, "fetch_k": 30, }, "retriever_llm_configuration": { @@ -300,7 +300,7 @@ } QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD = { - "summary": "multi-query reranker retriever with mmr", + "summary": "multi-query + re-ranking + mmr", "description": """ Typically used for complex user queries. Requires k and fetch_k in search_kwargs for mmr. @@ -322,9 +322,9 @@ "retriever_config": { "compressor_model_provider": "mixbread-ai", "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", - "top_k": 5, + "top_k": 7, "search_type": "similarity", - "search_kwargs": {"k": 20}, + "search_kwargs": {"k": 25}, "retriever_llm_configuration": { "name": "openai-main/gpt-3-5-turbo", "provider": "truefoundry", @@ -335,7 +335,7 @@ } QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD = { - "summary": "multi-query reranker retriever with similarity ", + "summary": "multi-query + re-ranking + similarity ", "description": """ Typically used for complex user queries. Requires k in search_kwargs for similarity search. @@ -343,3 +343,38 @@ "value": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY, } ####### + +QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": "Answer the question based only on the following context:\nContext: {context} \nQuestion: {question}", + "retriever_name": "contexual-compression-multi-query", + "retriever_config": { + "compressor_model_provider": "mixbread-ai", + "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", + "top_k": 7, + "search_type": "similarity_score_threshold", + "search_kwargs": {"score_threshold": 0.7}, + "retriever_llm_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + }, + "stream": False, +} + +QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD = { + "summary": "multi-query + re-ranking + threshold score", + "description": """ + Typically used for complex user queries. + Requires k in search_kwargs for similarity search. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE, +} +####### diff --git a/backend/modules/query_controllers/example/types.py b/backend/modules/query_controllers/example/types.py index acf2d46b..f922c471 100644 --- a/backend/modules/query_controllers/example/types.py +++ b/backend/modules/query_controllers/example/types.py @@ -116,7 +116,7 @@ class ExampleQueryInput(BaseModel): title="Collection name on which to search", ) - query: str = Field(title="Question to search for", max_length=1000) + query: str = Field(title="Question to search for") model_configuration: LLMConfig diff --git a/backend/modules/query_controllers/summary/controller.py b/backend/modules/query_controllers/summary/controller.py new file mode 100644 index 00000000..bf129c0f --- /dev/null +++ b/backend/modules/query_controllers/summary/controller.py @@ -0,0 +1,387 @@ +import asyncio +import json + +import async_timeout +from fastapi import Body, HTTPException +from fastapi.responses import StreamingResponse +from langchain.prompts import PromptTemplate +from langchain.retrievers import ContextualCompressionRetriever, MultiQueryRetriever +from langchain.schema.vectorstore import VectorStoreRetriever +from langchain_community.chat_models.ollama import ChatOllama +from langchain_core.output_parsers import StrOutputParser +from langchain_core.runnables import RunnableParallel, RunnablePassthrough +from langchain_openai.chat_models import ChatOpenAI +from truefoundry.langchain import TrueFoundryChat + +from backend.logger import logger +from backend.modules.embedder.embedder import get_embedder +from backend.modules.metadata_store.client import METADATA_STORE_CLIENT +from backend.modules.query_controllers.summary.payload import ( + QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD, + QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, + QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, + QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_PAYLOAD, + QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_MMR_PAYLOAD, + QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE_PAYLOAD, + QUERY_WITH_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD, + QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, + QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, + QUERY_WITH_VECTOR_STORE_RETRIEVER_MMR_PAYLOAD, + QUERY_WITH_VECTOR_STORE_RETRIEVER_PAYLOAD, + QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, +) +from backend.modules.query_controllers.summary.types import ( + GENERATION_TIMEOUT_SEC, + ExampleQueryInput, +) +from backend.modules.reranker import MxBaiReranker +from backend.modules.vector_db.client import VECTOR_STORE_CLIENT +from backend.server.decorators import post, query_controller +from backend.settings import settings + +EXAMPLES = { + "vector-store-similarity": QUERY_WITH_VECTOR_STORE_RETRIEVER_PAYLOAD, + "vector-store-similarity-threshold": QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, + "contexual-compression-multi-query-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, + "contextual-compression-multi-query-similarity-threshold": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, + # Keeping these for future use: + # "contexual-compression-similarity-threshold": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE_PAYLOAD, + # "vector-store-mmr": QUERY_WITH_VECTOR_STORE_RETRIEVER_MMR_PAYLOAD, + # "contexual-compression-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_PAYLOAD, + # "multi-query-mmr": QUERY_WITH_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD, +} + +# if settings.LOCAL: +# EXAMPLES.update( +# { +# "contexual-compression-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_PAYLOAD, +# "contexual-compression-multi-query-similarity": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, +# "multi-query-similarity": QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD, +# "multi-query-similarity-threshold": QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD, +# } +# ) + + +@query_controller("/summary-report") +class SummaryQueryController: + def _get_prompt_template(self, input_variables, template): + """ + Get the prompt template + """ + return PromptTemplate(input_variables=input_variables, template=template) + + def _format_docs(self, docs): + return "\n\n".join(doc.page_content for doc in docs) + + def _format_docs_for_stream(self, docs): + return [ + {"page_content": doc.page_content, "metadata": doc.metadata} for doc in docs + ] + + def _get_llm(self, model_configuration, stream=False): + """ + Get the LLM + """ + system = "You are a helpful assistant." + if model_configuration.provider == "openai": + logger.debug(f"Using OpenAI model {model_configuration.name}") + llm = ChatOpenAI( + model=model_configuration.name, + temperature=model_configuration.parameters.get("temperature", 0.1), + streaming=stream, + ) + elif model_configuration.provider == "ollama": + logger.debug(f"Using Ollama model {model_configuration.name}") + llm = ChatOllama( + base_url=settings.OLLAMA_URL, + model=( + model_configuration.name.split("/")[1] + if "/" in model_configuration.name + else model_configuration.name + ), + temperature=model_configuration.parameters.get("temperature", 0.1), + system=system, + ) + elif model_configuration.provider == "truefoundry": + logger.debug(f"Using TrueFoundry model {model_configuration.name}") + llm = TrueFoundryChat( + model=model_configuration.name, + model_parameters=model_configuration.parameters, + system_prompt=system, + ) + else: + logger.debug(f"Using TrueFoundry model {model_configuration.name}") + llm = TrueFoundryChat( + model=model_configuration.name, + model_parameters=model_configuration.parameters, + system_prompt=system, + ) + return llm + + async def _get_vector_store(self, collection_name: str): + """ + Get the vector store for the collection + """ + collection = METADATA_STORE_CLIENT.get_collection_by_name(collection_name) + + if collection is None: + raise HTTPException(status_code=404, detail="Collection not found") + + return VECTOR_STORE_CLIENT.get_vector_store( + collection_name=collection.name, + embeddings=get_embedder(collection.embedder_config), + ) + + def _get_vector_store_retriever(self, vector_store, retriever_config): + """ + Get the vector store retriever + """ + return VectorStoreRetriever( + vectorstore=vector_store, + search_type=retriever_config.search_type, + search_kwargs=retriever_config.search_kwargs, + ) + + def _get_contextual_compression_retriever(self, vector_store, retriever_config): + """ + Get the contextual compression retriever + """ + # Using mixbread-ai Reranker + if retriever_config.compressor_model_provider == "mixbread-ai": + retriever = self._get_vector_store_retriever(vector_store, retriever_config) + + compressor = MxBaiReranker( + model=retriever_config.compressor_model_name, + top_k=retriever_config.top_k, + ) + + compression_retriever = ContextualCompressionRetriever( + base_compressor=compressor, base_retriever=retriever + ) + + return compression_retriever + # Can add other rerankers too! + else: + raise HTTPException( + status_code=404, detail="Compressor model provider not found" + ) + + def _get_multi_query_retriever( + self, vector_store, retriever_config, retriever_type="vectorstore" + ): + """ + Get the multi query retriever + """ + if retriever_type == "vectorstore": + base_retriever = self._get_vector_store_retriever( + vector_store, retriever_config + ) + elif retriever_type == "contexual-compression": + base_retriever = self._get_contextual_compression_retriever( + vector_store, retriever_config + ) + + return MultiQueryRetriever.from_llm( + retriever=base_retriever, + llm=self._get_llm(retriever_config.retriever_llm_configuration), + ) + + async def _get_retriever(self, vector_store, retriever_name, retriever_config): + """ + Get the retriever + """ + if retriever_name == "vectorstore": + logger.debug( + f"Using VectorStoreRetriever with {retriever_config.search_type} search" + ) + retriever = self._get_vector_store_retriever(vector_store, retriever_config) + + elif retriever_name == "contexual-compression": + logger.debug( + f"Using ContextualCompressionRetriever with {retriever_config.search_type} search" + ) + retriever = self._get_contextual_compression_retriever( + vector_store, retriever_config + ) + + elif retriever_name == "multi-query": + logger.debug( + f"Using MultiQueryRetriever with {retriever_config.search_type} search" + ) + retriever = self._get_multi_query_retriever(vector_store, retriever_config) + + elif retriever_name == "contexual-compression-multi-query": + logger.debug( + f"Using MultiQueryRetriever with {retriever_config.search_type} search and retriever type as contexual-compression" + ) + retriever = self._get_multi_query_retriever( + vector_store, retriever_config, retriever_type="contexual-compression" + ) + + else: + raise HTTPException(status_code=404, detail="Retriever not found") + + return retriever + + async def _stream_answer(self, rag_chain, query): + async with async_timeout.timeout(GENERATION_TIMEOUT_SEC): + try: + async for chunk in rag_chain.astream(query): + if "question " in chunk: + # print("Question: ", chunk['question']) + yield json.dumps({"question": chunk["question"]}) + await asyncio.sleep(0.1) + elif "context" in chunk: + # print("Context: ", self._format_docs_for_stream(chunk['context'])) + yield json.dumps( + {"docs": self._format_docs_for_stream(chunk["context"])} + ) + await asyncio.sleep(0.1) + elif "answer" in chunk: + # print("Answer: ", chunk['answer']) + yield json.dumps({"answer": chunk["answer"]}) + await asyncio.sleep(0.1) + + yield json.dumps({"end": ""}) + except asyncio.TimeoutError: + raise HTTPException(status_code=504, detail="Stream timed out") + + @post("/answer") + async def answer( + self, + request: ExampleQueryInput = Body( + openapi_examples=EXAMPLES, + ), + ): + """ + Sample answer method to answer the question using the context from the collection + """ + try: + # Get the vector store + vector_store = await self._get_vector_store(request.collection_name) + + # Create the QA prompt templates + QA_PROMPT = self._get_prompt_template( + input_variables=["context", "question"], + template=request.prompt_template, + ) + + # Get the LLM + llm = self._get_llm(request.model_configuration, request.stream) + + # get retriever + retriever = await self._get_retriever( + vector_store=vector_store, + retriever_name=request.retriever_name, + retriever_config=request.retriever_config, + ) + + # Using LCEL + rag_chain_from_docs = ( + RunnablePassthrough.assign( + context=(lambda x: self._format_docs(x["context"])) + ) + | QA_PROMPT + | llm + | StrOutputParser() + ) + + rag_chain_with_source = RunnableParallel( + {"context": retriever, "question": RunnablePassthrough()} + ).assign(answer=rag_chain_from_docs) + + if request.stream: + return StreamingResponse( + self._stream_answer(rag_chain_with_source, request.query), + media_type="text/event-stream", + ) + + else: + outputs = await rag_chain_with_source.ainvoke(request.query) + + # Intermediate testing + # Just the retriever + # setup_and_retrieval = RunnableParallel({"context": retriever, "question": RunnablePassthrough()}) + # outputs = await setup_and_retrieval.ainvoke(request.query) + # print(outputs) + + # Retriver and QA + # outputs = await (setup_and_retrieval | QA_PROMPT).ainvoke(request.query) + # print(outputs) + + # Retriver, QA and LLM + # outputs = await (setup_and_retrieval | QA_PROMPT | llm).ainvoke(request.query) + # print(outputs) + + SUMMARY_PROMPT = "You are an AI assistant specialising in summarizing documents finance, insurance and private equity. Given a list of question and answers, your task is to provide a detailed one pager summary report. Summary: {context}" + + # Get the summary + summary_rag_chain = ( + RunnablePassthrough.assign( + context=lambda x: x["answer"], + ) + | PromptTemplate( + input_variables=["context"], + template=SUMMARY_PROMPT, + ) + | llm + | StrOutputParser() + ) + + summary = await summary_rag_chain.ainvoke(outputs) + + answer = ( + outputs["answer"] + "\n\n**Summary:**\n" + summary + if ("Summary" or "summary") not in summary + else outputs["answer"] + "\n\n\n" + summary + ) + return { + "answer": answer, + "docs": outputs["context"] if outputs["context"] else [], + } + + except HTTPException as exp: + raise exp + except Exception as exp: + logger.exception(exp) + raise HTTPException(status_code=500, detail=str(exp)) + + +####### +# Streaming Client + +# import httpx +# from httpx import Timeout + +# from backend.modules.query_controllers.summary.types import ExampleQueryInput + +# payload = { +# "collection_name": "pstest", +# "query": "What are the features of Diners club black metal edition?", +# "model_configuration": { +# "name": "openai-devtest/gpt-3-5-turbo", +# "parameters": { +# "temperature": 0.1 +# }, +# "provider": "truefoundry" +# }, +# "prompt_template": "Answer the question based only on the following context:\nContext: {context} \nQuestion: {question}", +# "retriever_name": "vectorstore", +# "retriever_config": { +# "search_type": "similarity", +# "search_kwargs": { +# "k": 20 +# }, +# "filter": {} +# }, +# "stream": True +# } + +# data = ExampleQueryInput(**payload).dict() +# ENDPOINT_URL = 'http://localhost:8000/retrievers/example-app/answer' + + +# with httpx.stream('POST', ENDPOINT_URL, json=data, timeout=Timeout(5.0*60)) as r: +# for chunk in r.iter_text(): +# print(chunk) +####### diff --git a/backend/modules/query_controllers/summary/payload.py b/backend/modules/query_controllers/summary/payload.py new file mode 100644 index 00000000..8083866a --- /dev/null +++ b/backend/modules/query_controllers/summary/payload.py @@ -0,0 +1,389 @@ +SUMMARY_PROMPT_TEMPLATE = """You are an AI assistant specialising in finance, insurance and private equity. Your role is to provide detailed answers of 60-80 words, drawing directly from the context. +In your responses: +- Only use information found in the listed sources +- Do not generate speculative answers or information not supported by the sources +- Each question should be answered separately in question and answer pair format. +Context: {context} +Question: {question} +""" + +QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "vectorstore", + "retriever_config": {"search_type": "similarity", "search_kwargs": {"k": 5}}, + "stream": False, +} + +QUERY_WITH_VECTOR_STORE_RETRIEVER_PAYLOAD = { + "summary": "search with similarity", + "description": """ + Requires k in search_kwargs for similarity search. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY, +} +####### + +QUERY_WITH_VECTOR_STORE_RETRIEVER_MMR = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "vectorstore", + "retriever_config": { + "search_type": "mmr", + "search_kwargs": { + "k": 5, + "fetch_k": 7, + }, + }, + "stream": False, +} + +QUERY_WITH_VECTOR_STORE_RETRIEVER_MMR_PAYLOAD = { + "summary": "search with mmr", + "description": """ + Requires k and fetch_k in search_kwargs for mmr support depends on vector db. + search_type can either be similarity or mmr or similarity_score_threshold""", + "value": QUERY_WITH_VECTOR_STORE_RETRIEVER_MMR, +} +####### + +QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY_SCORE = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "vectorstore", + "retriever_config": { + "search_type": "similarity_score_threshold", + "search_kwargs": {"score_threshold": 0.7}, + }, + "stream": False, +} + +QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY_SCORE_PAYLOAD = { + "summary": "search with threshold score", + "description": """ + Requires score_threshold float (0~1) in search kwargs. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_VECTOR_STORE_RETRIEVER_SIMILARITY_SCORE, +} +####### + +QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "contexual-compression", + "retriever_config": { + "compressor_model_provider": "mixbread-ai", + "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", + "top_k": 7, + "search_type": "similarity", + "search_kwargs": {"k": 25}, + }, + "stream": False, +} + +QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_PAYLOAD = { + "summary": "similariy search + re-ranking", + "description": """ + Requires k in search_kwargs for similarity search. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER, +} +##### + + +QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_MMR = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "contexual-compression", + "retriever_config": { + "compressor_model_provider": "mixbread-ai", + "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", + "top_k": 7, + "search_type": "mmr", + "search_kwargs": { + "k": 25, + "fetch_k": 30, + }, + }, + "stream": False, +} + +QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_MMR_PAYLOAD = { + "summary": "mmr + re-ranking", + "description": """ + Requires k and fetch_k in search kwargs for mmr. + search_type can either be similarity or mmr or similarity_score_threshold. + Currently only support for mixedbread-ai/mxbai-rerank-xsmall-v1 reranker is added.""", + "value": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_MMR, +} + +##### + + +QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "contexual-compression", + "retriever_config": { + "compressor_model_provider": "mixbread-ai", + "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", + "top_k": 7, + "search_type": "similarity_score_threshold", + "search_kwargs": {"score_threshold": 0.7}, + }, + "stream": False, +} + +QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE_PAYLOAD = { + "summary": "threshold score + re-ranking", + "description": """ + Requires score_threshold float (0~1) in search kwargs for similarity search. + search_type can either be similarity or mmr or similarity_score_threshold. + Currently only support for mixedbread-ai/mxbai-rerank-xsmall-v1 reranker is added""", + "value": QUERY_WITH_CONTEXTUAL_COMPRESSION_RETRIEVER_SEARCH_TYPE_SIMILARITY_WITH_SCORE, +} + +##### + +QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "multi-query", + "retriever_config": { + "search_type": "similarity", + "search_kwargs": {"k": 5}, + "retriever_llm_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.9}, + }, + }, + "stream": False, +} + +QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD = { + "summary": "multi-query + similarity search", + "description": """ + Typically used for complex user queries. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY, +} +####### + + +QUERY_WITH_MULTI_QUERY_RETRIEVER_MMR = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "multi-query", + "retriever_config": { + "search_type": "mmr", + "search_kwargs": { + "k": 5, + "fetch_k": 10, + }, + "retriever_llm_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.9}, + }, + }, + "stream": False, +} + +QUERY_WITH_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD = { + "summary": "multi-query + mmr", + "description": """ + Requires k and fetch_k in search_kwargs for mmr. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_MULTI_QUERY_RETRIEVER_MMR, +} +####### + +QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "multi-query", + "retriever_config": { + "search_type": "similarity_score_threshold", + "search_kwargs": {"score_threshold": 0.7}, + "retriever_llm_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.9}, + }, + }, + "stream": False, +} + +QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD = { + "summary": "multi-query + threshold score", + "description": """ + Typically used for complex user queries. + Requires score_threshold float (0~1) in search kwargs. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE, +} +####### + + +QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_MMR = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "contexual-compression-multi-query", + "retriever_config": { + "compressor_model_provider": "mixbread-ai", + "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", + "top_k": 7, + "search_type": "mmr", + "search_kwargs": { + "k": 25, + "fetch_k": 30, + }, + "retriever_llm_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.9}, + }, + }, + "stream": False, +} + +QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_MMR_PAYLOAD = { + "summary": "multi-query + re-ranking + mmr", + "description": """ + Typically used for complex user queries. + Requires k and fetch_k in search_kwargs for mmr. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_MMR, +} +####### + +QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "contexual-compression-multi-query", + "retriever_config": { + "compressor_model_provider": "mixbread-ai", + "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", + "top_k": 7, + "search_type": "similarity", + "search_kwargs": {"k": 25}, + "retriever_llm_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + }, + "stream": False, +} + +QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_PAYLOAD = { + "summary": "multi-query + re-ranking + similarity ", + "description": """ + Typically used for complex user queries. + Requires k in search_kwargs for similarity search. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY, +} +####### + +QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE = { + "collection_name": "creditcard", + "query": "Explain in detail different categories of credit cards", + "model_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + "prompt_template": SUMMARY_PROMPT_TEMPLATE, + "retriever_name": "contexual-compression-multi-query", + "retriever_config": { + "compressor_model_provider": "mixbread-ai", + "compressor_model_name": "mixedbread-ai/mxbai-rerank-xsmall-v1", + "top_k": 7, + "search_type": "similarity_score_threshold", + "search_kwargs": {"score_threshold": 0.7}, + "retriever_llm_configuration": { + "name": "openai-main/gpt-3-5-turbo", + "provider": "truefoundry", + "parameters": {"temperature": 0.1}, + }, + }, + "stream": False, +} + +QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE_PAYLOAD = { + "summary": "multi-query + re-ranking + threshold score", + "description": """ + Typically used for complex user queries. + Requires k in search_kwargs for similarity search. + search_type can either be similarity or mmr or similarity_score_threshold.""", + "value": QUERY_WITH_CONTEXTUAL_COMPRESSION_MULTI_QUERY_RETRIEVER_SIMILARITY_SCORE, +} +####### diff --git a/backend/modules/query_controllers/summary/types.py b/backend/modules/query_controllers/summary/types.py new file mode 100644 index 00000000..f922c471 --- /dev/null +++ b/backend/modules/query_controllers/summary/types.py @@ -0,0 +1,178 @@ +from typing import Any, ClassVar, Collection, Dict, Literal, Optional + +from pydantic import BaseModel, Field, root_validator, validator +from qdrant_client.models import Filter as QdrantFilter + +from backend.types import LLMConfig + +GENERATION_TIMEOUT_SEC = 60.0 * 5 + + +class VectorStoreRetrieverConfig(BaseModel): + """ + Configuration for VectorStore Retriever + """ + + search_type: str = Field( + default="similarity", + title="""Defines the type of search that the Retriever should perform. Can be 'similarity' (default), 'mmr', or 'similarity_score_threshold'. + - "similarity": Retrieve the top k most similar documents to the query., + - "mmr": Retrieve the top k most similar documents to the query and then rerank them using Maximal Marginal Relevance (MMR)., + - "similarity_score_threshold": Retrieve all documents with similarity score greater than a threshold. + """, + ) + + search_kwargs: dict = Field(default_factory=dict) + + filter: Optional[dict] = Field( + default_factory=dict, + title="""Filter by document metadata""", + ) + + allowed_search_types: ClassVar[Collection[str]] = ( + "similarity", + "similarity_score_threshold", + "mmr", + ) + + @root_validator + def validate_search_type(cls, values: Dict) -> Dict: + """Validate search type.""" + search_type = values.get("search_type") + + assert ( + search_type in cls.allowed_search_types + ), f"search_type of {search_type} not allowed. Valid values are: {cls.allowed_search_types}" + + search_kwargs = values.get("search_kwargs") + + if search_type == "similarity": + assert "k" in search_kwargs, "k is required for similarity search" + + elif search_type == "mmr": + assert "k" in search_kwargs, "k is required in search_kwargs for mmr search" + assert ( + "fetch_k" in search_kwargs + ), "fetch_k is required in search_kwargs for mmr search" + + elif search_type == "similarity_score_threshold": + assert ( + "score_threshold" in search_kwargs + ), "score_threshold with a float value(0~1) is required in search_kwargs for similarity_score_threshold search" + + filters = values.get("filter") + if filters: + search_kwargs["filter"] = QdrantFilter.parse_obj(filters) + return values + + +class MultiQueryRetrieverConfig(VectorStoreRetrieverConfig): + retriever_llm_configuration: LLMConfig = Field( + title="LLM configuration for the retriever", + ) + + +class ContextualCompressionRetrieverConfig(VectorStoreRetrieverConfig): + compressor_model_provider: str = Field( + title="provider of the compressor model", + ) + + compressor_model_name: str = Field( + title="model name of the compressor", + ) + + top_k: int = Field( + title="Top K docs to collect post compression", + ) + + allowed_compressor_model_providers: ClassVar[Collection[str]] = ("mixbread-ai",) + + @validator("compressor_model_provider") + def validate_retriever_type(cls, value) -> Dict: + assert ( + value in cls.allowed_compressor_model_providers + ), f"Compressor model of {value} not allowed. Valid values are: {cls.allowed_compressor_model_providers}" + return value + + +class ContextualCompressionMultiQueryRetrieverConfig( + ContextualCompressionRetrieverConfig, MultiQueryRetrieverConfig +): + pass + + +class LordOfRetrievers(ContextualCompressionRetrieverConfig, MultiQueryRetrieverConfig): + pass + + +class ExampleQueryInput(BaseModel): + """ + Model for Query input. + Requires a collection name, retriever configuration, query, LLM configuration and prompt template. + """ + + collection_name: str = Field( + default=None, + title="Collection name on which to search", + ) + + query: str = Field(title="Question to search for") + + model_configuration: LLMConfig + + prompt_template: str = Field( + title="Prompt Template to use for generating answer to the question using the context", + ) + + retriever_name: str = Field( + title="Retriever name", + ) + + retriever_config: Dict[str, Any] = Field( + title="Retriever configuration", + ) + + allowed_retriever_types: ClassVar[Collection[str]] = ( + "vectorstore", + "multi-query", + "contexual-compression", + "contexual-compression-multi-query", + "lord-of-the-retrievers", + ) + + stream: Optional[bool] = Field(title="Stream the results", default=False) + + @root_validator() + def validate_retriever_type(cls, values: Dict) -> Dict: + retriever_name = values.get("retriever_name") + + assert ( + retriever_name in cls.allowed_retriever_types + ), f"retriever of {retriever_name} not allowed. Valid values are: {cls.allowed_retriever_types}" + + if retriever_name == "vectorstore": + values["retriever_config"] = VectorStoreRetrieverConfig( + **values.get("retriever_config") + ) + + elif retriever_name == "multi-query": + values["retriever_config"] = MultiQueryRetrieverConfig( + **values.get("retriever_config") + ) + + elif retriever_name == "contexual-compression": + values["retriever_config"] = ContextualCompressionRetrieverConfig( + **values.get("retriever_config") + ) + + elif retriever_name == "contexual-compression-multi-query": + values["retriever_config"] = ContextualCompressionMultiQueryRetrieverConfig( + **values.get("retriever_config") + ) + + elif retriever_name == "lord-of-the-retrievers": + values["retriever_config"] = LordOfRetrievers( + **values.get("retriever_config") + ) + + return values diff --git a/backend/settings.py b/backend/settings.py index da7b533d..79f955f3 100644 --- a/backend/settings.py +++ b/backend/settings.py @@ -63,5 +63,7 @@ class Settings(BaseSettings): except Exception as e: raise ValueError(f"METADATA_STORE_CONFIG is invalid: {e}") + TF_PARSER = os.getenv("TF_PARSER", "") + settings = Settings() diff --git a/local.metadata.yaml b/local.metadata.yaml index a02b7077..17762c96 100644 --- a/local.metadata.yaml +++ b/local.metadata.yaml @@ -1,17 +1,17 @@ -collection_name: creditcard +collection_name: bradysmall data_source: type: localdir - uri: sample-data/creditcards + uri: /Users/prathamesh/Desktop/brady parser_config: chunk_size: 1000 chunk_overlap: 20 parser_map: - ".md": MarkdownParser + ".pdf": TfParser embedder_config: - provider: mixedbread - config: - model: mixedbread-ai/mxbai-embed-large-v1 - # For Truefoundry Embeddings - # provider: truefoundry + # provider: mixedbread # config: - # model: "openai-main/text-embedding-ada-002" + # model: mixedbread-ai/mxbai-embed-large-v1 + # For Truefoundry Embeddings + provider: truefoundry + config: + model: "openai-main/text-embedding-ada-002"