rag update

app.py

@@ -11,8 +11,9 @@ from models import (
     UrgencyPrediction
 )
 from multi_task_model_class import MultiTaskModel
-from rag import generate_answer
+
 from memory import get_conversation, add_message
+from service.rag_service import generate_answer
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -127,15 +128,6 @@ def classify_text(text: str, threshold: float = 0.5):
     return categories, urgency
 
 
-def retrieve_documents(query: str):
-    """
-    Retrieve relevant documents for RAG.
-    """
-    return [
-        "Restarting the router fixes most connectivity issues.",
-        "Check for planned ISP maintenance.",
-        "Verify cables are securely connected."
-    ]
 
 
 @app.get("/")
@@ -160,8 +152,7 @@ def query_endpoint(req: QueryRequest):
     categories, urgency = classify_text(req.query)
 
     # RAG
-    docs = retrieve_documents(req.query)
-    answer = generate_answer(req.query, docs, history)
+    answer = generate_answer(req.query,history)
 
     # Update conversation memory
     add_message(req.user_id, "user", req.query)

service/data_loader_service.py

@@ -0,0 +1,22 @@
+import csv
+from pathlib import Path
+
+class CSVDataLoader:
+    def __init__(self, file_path: str):
+        self.file_path = Path(file_path)
+
+    def load_qa_pairs(self) -> list[str]:
+        documents = []
+
+        with open(self.file_path, newline="", encoding="utf-8") as f:
+            reader = csv.DictReader(f)
+
+            for row in reader:
+                question = row.get("question", "").strip()
+                answer = row.get("answer", "").strip()
+
+                if question and answer:
+                    doc = f"Question: {question}\nAnswer: {answer}"
+                    documents.append(doc)
+
+        return documents

service/embedded_service.py

@@ -0,0 +1,8 @@
+from sentence_transformers import SentenceTransformer
+
+class EmbeddingService:
+    def __init__(self):
+        self.model = SentenceTransformer("all-MiniLM-L6-v2")
+
+    def embed(self, texts: list[str]):
+        return self.model.encode(texts, normalize_embeddings=True)

service/llm_service.py

@@ -0,0 +1,26 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+class LLMService:
+    def __init__(self):
+        model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+
+    def generate(self, prompt: str) -> str:
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+
+        output = self.model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.3,
+            top_p=0.9,
+            do_sample=True
+        )
+
+        return self.tokenizer.decode(output[0], skip_special_tokens=True)

service/memory_service.py

@@ -0,0 +1,7 @@
+memory_store = {}
+
+def get_memory(session_id: str) -> str:
+    return memory_store.get(session_id, "")
+
+def save_memory(session_id: str, summary: str):
+    memory_store[session_id] = summary

service/rag_service.py

@@ -0,0 +1,43 @@
+
+from service.embedded_service import EmbeddingService
+from service.llm_service import LLMService
+from service.memory_service import get_memory
+from service.vector_store_service import VectorStoreService
+from service.data_loader_service import CSVDataLoader
+
+# Initialize once
+embedder = EmbeddingService()
+llm = LLMService()
+
+# Load documents from CSV
+loader = CSVDataLoader("final_data_set(in).csv")
+documents = loader.load_qa_pairs()
+
+doc_embeddings = embedder.embed(documents)
+vector_store = VectorStoreService(doc_embeddings, documents)
+
+def generate_answer(question: str, session_id: str | None):
+    query_embedding = embedder.embed([question])[0]
+    context_chunks = vector_store.search(query_embedding, k=3)
+
+    memory = get_memory(session_id) if session_id else ""
+
+    prompt = f"""
+<|system|>
+You are a helpful assistant. Answer ONLY using the provided context.
+If the answer is not in the context, say "I don't know".
+
+Conversation memory:
+{memory}
+
+<|user|>
+Context:
+{chr(10).join(context_chunks)}
+
+Question:
+{question}
+
+<|assistant|>
+"""
+
+    return llm.generate(prompt)

service/vector_store_service.py

@@ -0,0 +1,16 @@
+import faiss
+import numpy as np
+
+class VectorStoreService:
+    def __init__(self, embeddings: list, documents: list):
+        self.documents = documents
+        dim = len(embeddings[0])
+
+        self.index = faiss.IndexFlatL2(dim)
+        self.index.add(np.array(embeddings))
+
+    def search(self, query_embedding, k=3):
+        distances, indices = self.index.search(
+            np.array([query_embedding]), k
+        )
+        return [self.documents[i] for i in indices[0]]