From 0f8469c03b44fe082b609ef0f9676a66cbd94d41 Mon Sep 17 00:00:00 2001
From: KevKibe <kchegz234@gmail.com>
Date: Tue, 30 Apr 2024 17:34:40 +0300
Subject: [PATCH] add utils module with Query result pydantic model and prompt
 configuration

---
 src/utils/__init__.py       |  0
 src/utils/config.py         | 14 ++++++++++++++
 src/utils/response_model.py | 38 +++++++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+)
 create mode 100644 src/utils/__init__.py
 create mode 100644 src/utils/config.py
 create mode 100644 src/utils/response_model.py

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'
+#         }
+#     ]
+# }
+