Update main.py

main.py

@@ -1,42 +1,48 @@
 import os
-import json
+import chromadb
+from chromadb.utils import embedding_functions
 from groq import Groq
 
 class AICore:
     def __init__(self, api_key):
         self.client = Groq(api_key=api_key)
-        self.model = "qwen-2-vl-7b-instruct"  # Sanjay bhai, vision ke liye best hai
-        self.memory_file = "memory.json"
-        self._load_memory()
+        self.chroma_client = chromadb.PersistentClient(path="./chroma_db")
+        self.embed_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
+        self.collection = self.chroma_client.get_or_create_collection(name="aumcore_memory", embedding_function=self.embed_fn)
 
-    def _load_memory(self):
-        if os.path.exists(self.memory_file):
-            with open(self.memory_file, 'r') as f:
-                self.memory = json.load(f)
-        else:
-            self.memory = []
+    def get_response(self, user_input, user_name="Sanjay"):
+        # Purani yaadein fetch karna (Fact-finding only)
+        results = self.collection.query(query_texts=[user_input], n_results=2)
+        memory_context = "\n".join(results['documents'][0]) if results['documents'] else "No past history."
 
-    def save_to_memory(self, user_input, ai_response):
-        self.memory.append({"user": user_input, "bot": ai_response})
-        with open(self.memory_file, 'w') as f:
-            json.dump(self.memory[-10:], f) # Last 10 chats yaad rakhega
-
-    def get_response(self, text, image_path=None):
-        try:
-            # Simple text + vision prompt setup
-            messages = [{"role": "user", "content": text}]
-            # Future: Image processing logic can be expanded here
-            
-            completion = self.client.chat.completions.create(
-                model=self.model,
-                messages=messages,
-                temperature=0.7
-            )
-            response = completion.choices[0].message.content
-            self.save_to_memory(text, response)
-            return response
-        except Exception as e:
-            return f"Error: {str(e)}"
-
-if __name__ == "__main__":
-    print("AICore is ready!")
+        # Strict Persona Enforcement (No basic replies allowed)
+        system_prompt = f"""
+        Role: Senior Software Architect (AumCore AI).
+        User: {user_name}.
+        Language: Strictly 60% English & 40% Hindi.
+        Rule: If asking for code, provide a high-level Python script with try-except blocks. 
+        Style: Do NOT be generic. Be professional and helpful.
+        
+        Information from Memory: {memory_context}
+        """
+        
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": "colab me gpu check karne ka code do"},
+            {"role": "assistant", "content": f"{user_name} bhai, Colab mein GPU visibility confirm karne ke liye ye standard script use karein.\n\n```python\nimport torch\ntry:\n    status = torch.cuda.is_available()\n    print(f'GPU Available: {{status}}')\n    if status: print(f'Device: {{torch.cuda.get_device_name(0)}}')\nexcept Exception as e:\n    print(f'Error identifying GPU: {{e}}')\n```"},
+            {"role": "user", "content": user_input}
+        ]
+        
+        completion = self.client.chat.completions.create(
+            model="llama-3.3-70b-versatile",
+            messages=messages,
+            temperature=0.1 # Low temperature for consistency
+        )
+        response = completion.choices[0].message.content
+        
+        # Saving new memory
+        self.collection.add(
+            documents=[f"User: {user_input} | AI: {response[:100]}"],
+            ids=[str(os.urandom(4).hex())]
+        )
+        return response