updated agent

agent.py

@@ -1,41 +1,32 @@
+# agent.py — Gemini-based Coding Agent with PDF/.py RAG + Memory (.env support)
 import os
 import fitz
 import faiss
-import torch
 import sqlite3
 import numpy as np
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import google.generativeai as genai
+from dotenv import load_dotenv
 from sentence_transformers import SentenceTransformer
-from torch.quantization import quantize_dynamic
 
 class CodingAgent:
     def __init__(self):
-        # Load TinyLlama (chat-optimized, CPU-compatible)
-        model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        self.model = AutoModelForCausalLM.from_pretrained(model_id)
+        # Load .env file
+        load_dotenv()
 
-        # Apply dynamic quantization to speed up inference on CPU
-        self.model = quantize_dynamic(
-            self.model, {torch.nn.Linear}, dtype=torch.qint8
-        )
+        # Setup Gemini LLM
+        self.api_key = os.getenv("GEMINI_API_KEY")
+        if not self.api_key:
+            raise ValueError("GEMINI_API_KEY not found in environment or .env file.")
 
-        # Text-generation pipeline: reduced token output and deterministic
-        self.llm = pipeline(
-            "text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
-            max_new_tokens=128,
-            do_sample=False,
-            device=-1
-        )
+        genai.configure(api_key=self.api_key)
+        self.model = genai.GenerativeModel("gemini-1.5-flash")
 
         # Embedding model + FAISS index
         self.embedder = SentenceTransformer("all-MiniLM-L6-v2")
         self.index = faiss.IndexFlatL2(384)
         self.docs = []
 
-        # SQLite for session memory
+        # SQLite memory
         self.conn = sqlite3.connect("memory.db", check_same_thread=False)
         self.conn.execute(
             """CREATE TABLE IF NOT EXISTS memory (id INTEGER PRIMARY KEY, query TEXT, response TEXT)"""
@@ -70,7 +61,7 @@ class CodingAgent:
         self.docs.extend(chunks)
         return f"Added {len(chunks)} chunks."
 
-    def retrieve_context(self, query, top_k=1):
+    def retrieve_context(self, query, top_k=2):
         if self.index.ntotal == 0:
             return ""
         query_emb = self.embed_chunks([query])[0]
@@ -78,7 +69,7 @@ class CodingAgent:
         return "\n\n".join(self.docs[i] for i in I[0])
 
     def answer(self, query):
-        # Check memory
+        # Check memory first
         cursor = self.conn.execute(
             "SELECT response FROM memory WHERE query = ?", (query,)
         )
@@ -88,19 +79,21 @@ class CodingAgent:
 
         context = self.retrieve_context(query)
         prompt = (
-            f"### User:\n{query}\n\n### Context:\n{context}\n\n### Assistant:"
+            f"You are a helpful coding assistant.\n\n"
+            f"Context (from uploaded docs):\n{context}\n\n"
+            f"User question: {query}\n\n"
+            f"Answer with code or explanation where needed."
         )
-        # Generate response
-        output = self.llm(prompt)[0]['generated_text']
-        response = output.split("### Assistant:")[-1].strip()
 
-        # Save to memory
+        response = self.model.generate_content(prompt)
+        answer = response.text.strip()
+
         self.conn.execute(
             "INSERT INTO memory (query, response) VALUES (?, ?)",
-            (query, response)
+            (query, answer)
         )
         self.conn.commit()
-        return response
+        return answer
 
     def clear_context(self):
         self.conn.execute("DELETE FROM memory")