add llm generation

utils/llm_generation.py

@@ -0,0 +1,84 @@
+from langchain.prompts import (
+    ChatPromptTemplate,
+    SystemMessagePromptTemplate,
+    HumanMessagePromptTemplate,
+)
+from langchain_community.chat_models import ChatOpenAI
+from langchain.chains import LLMChain
+from langchain.schema import Document
+from typing import List
+import json
+
+
+json_schema = {
+    "name": "generate_answer",
+    "strict": True,
+    "schema": {
+        "type": "object",
+        "properties": {
+            "answer": {"type": "string"},
+        },
+        "required": ["answer"],
+        "additionalProperties": False,
+    },
+}
+
+
+class LLMGeneration:
+    def __init__(self, llm_model_name="gpt-4o-mini"):
+        self.llm_model_name = llm_model_name
+        self.llm = ChatOpenAI(
+            model_name=self.llm_model_name,
+            temperature=0.1,
+            model_kwargs={
+                "response_format": {
+                    "type": "json_schema",
+                    "json_schema": json_schema,
+                }
+            },
+        )
+
+        self.create_system_message()
+
+    def create_system_message(self):
+        # System message for the chain
+        system_message = (
+            "You are a helpful assistant. Use ONLY the provided context to answer the user's question. "
+            'If you cannot find the answer in the context, return answer key as  "Data Not Available".'
+        )
+        self.system_message_prompt = SystemMessagePromptTemplate.from_template(
+            system_message
+        )
+
+    def create_human_message_prompt(self, query: str, docs: List[Document]):
+
+        # Prepare the context from the retrieved chunks
+        context = "\n\n".join([doc.page_content for doc in docs])
+
+        # Human message, instructing the assistant to use the context and produce a structured JSON answer
+        human_message = f"""Given the following context:
+
+        {context}
+
+        Answer the following question in a structured JSON format with a key called "answer". If you do not find the answer in the provided context, return "answer" as "No answer".
+
+        Question: {query}"""
+
+        return HumanMessagePromptTemplate.from_template(human_message)
+
+    def generate_answer(self, query: str, docs: List[Document]):
+
+        # Create the prompt template
+        prompt = ChatPromptTemplate.from_messages(
+            [self.system_message_prompt, self.create_human_message_prompt(query, docs)]
+        )
+
+        # Create and run the chain with the hypothetical gpt-40-mini model
+        chain = LLMChain(
+            llm=self.llm,
+            prompt=prompt,
+        )
+        result = chain.invoke({})
+
+        print(f"Query: {query} Result: {result}")
+        return result.get("answer", "")