diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/utils/config.py b/src/utils/config.py new file mode 100644 index 0000000..c5ecac3 --- /dev/null +++ b/src/utils/config.py @@ -0,0 +1,14 @@ +class Config: + template_str = """ + You are very helpful assistant for question answering tasks. Use the pieces of retrieved context to answer question given. If you do not know + the answer, Just say that you do not know the answer instead of making up an answer. + + Retrieved context: {context} + Query: {query} + format instructions: {format_instructions} + + """ + + default_google_model = "gemini-pro" + default_openai_model = "gpt-3.5-turbo-0125" + default_cohere_model = "command" diff --git a/src/utils/response_model.py b/src/utils/response_model.py new file mode 100644 index 0000000..fee8440 --- /dev/null +++ b/src/utils/response_model.py @@ -0,0 +1,38 @@ +from typing import List, Union +from pydantic import BaseModel, Field + +class Document(BaseModel): + page_content: str = Field(..., description="The content of the page from the source document.") + source: Union[float, int] = Field(..., description="The page number of the page_content in the document") + title: str = Field(..., description="The title or URL of the source document.") + + +class QueryResult(BaseModel): + query: str = Field(..., description="The query that was submitted.") + result: str = Field(..., description="The result of the query, including any retrieved information.") + page: Union[float, int] = Field(..., description="The page number of the final result of the query.") + source_documents: List[Document] = Field(..., description="A list of source documents related to the query.") + +# Example usage: +# data = { +# 'query': 'how did RAG come up?', +# 'result': 'RAG came up as a language model that is more strongly grounded in ' +# 'than BART and has been effective in Jeopardy question generation.\n' +# '\n' +# 'Sources:\n' +# '- https://arxiv.org/pdf/2005.11401.pdf (page 5.0)', +# 'page': 5.0, # A +# 'source_documents': [ +# { +# 'page_content': 'page-content-where-the-response-is-from.\n10', +# 'source': 9.0, +# 'title': 'document-title' +# }, +# { +# 'page_content': 'page-content-where-the-response-is-from.\n17', +# 'source': 5.0, +# 'title': 'document-title' +# } +# ] +# } +