refactor(ai): Upgrade langchain to v0.2.x syntax

core/support_agent.py

@@ -1,96 +1,138 @@
 import traceback
 from typing import Dict, Any, List
-from llama_cpp import Llama
 
-from langchain_core.language_models.llms import LLM
-from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain
-from langchain.memory import ConversationBufferMemory
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from llama_cpp import Llama
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
 from langchain_community.vectorstores import Chroma
-from langchain_core.prompts import PromptTemplate
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.language_models.llms import LLM
 
+# A custom wrapper to make llama_cpp compatible with LangChain's LLM interface
 class LlamaLangChain(LLM):
     llama_instance: Llama
+
     @property
-    def _llm_type(self) -> str: return "custom"
+    def _llm_type(self) -> str:
+        return "custom-llama-cpp"
+
     def _call(self, prompt: str, stop: List[str] | None = None, **kwargs) -> str:
-        response = self.llama_instance(prompt, max_tokens=512, stop=stop, stream=False, echo=False)
-        return response["choices"][0]["text"]
+        # Some LLMs may return conversational artifacts; we strip them here.
+        # This is a robust way to ensure a clean response.
+        unwanted_starters = ["Answer:", "Direct Answer:", "Assistant:"]
+        try:
+            response = self.llama_instance(prompt, max_tokens=512, stop=stop, stream=False, echo=False)
+            text = response["choices"][0]["text"].strip()
+            for starter in unwanted_starters:
+                if text.lower().startswith(starter.lower()):
+                    text = text[len(starter):].strip()
+            return text
+        except Exception as e:
+            print(f"ERROR during LLM call: {e}")
+            return "Error generating response from the model."
+
     async def _acall(self, prompt: str, stop: List[str] | None = None, **kwargs) -> str:
+        # Simple async wrapper around the synchronous call
         return self._call(prompt, stop, **kwargs)
 
-def format_docs(docs):
+# Helper function to format retrieved documents
+def _format_docs_for_context(docs: List[Any]) -> str:
     return "\n\n".join(doc.page_content for doc in docs)
 
 class SupportAgent:
+    """
+    Modern (LangChain v0.2.x) AI agent using a RAG pipeline with LCEL.
+    This version replaces the deprecated ConversationalRetrievalChain.
+    """
     def __init__(self, llm_instance: Llama, embedding_path: str, db_path: str):
-        print("--- Initializing Support Agent (Final Version) ---")
-        if llm_instance is None: raise ValueError("SupportAgent received an invalid LLM instance.")
+        print("--- Initializing Support Agent (LangChain v0.2.x Modern Version) ---")
+        if llm_instance is None:
+            raise ValueError("SupportAgent received an invalid LLM instance.")
+        
         self.langchain_llm_wrapper = LlamaLangChain(llama_instance=llm_instance)
-        self.embeddings = HuggingFaceEmbeddings(model_name=embedding_path)
-        self.vector_store = Chroma(persist_directory=db_path, embedding_function=self.embeddings)
-        self.conversations: Dict[str, ConversationBufferMemory] = {}
-        print("✅ Agent and core components initialized successfully.")
+        
+        try:
+            print(f"   - Loading embeddings from: {embedding_path}")
+            self.embeddings = HuggingFaceEmbeddings(model_name=embedding_path, model_kwargs={'device': 'cpu'})
+            
+            print(f"   - Connecting to Vector DB at: {db_path}")
+            self.vector_store = Chroma(persist_directory=db_path, embedding_function=self.embeddings)
+            self.retriever = self.vector_store.as_retriever(search_kwargs={"k": 3})
+
+            # The memory is no longer part of the chain itself in modern LCEL
+            self.conversations: Dict[str, List[tuple]] = {}
 
-    def _get_or_create_memory(self, conversation_id: str) -> ConversationBufferMemory:
-        if conversation_id not in self.conversations:
-            self.conversations[conversation_id] = ConversationBufferMemory(memory_key="chat_history", return_messages=True, input_key="question", output_key='answer')
-        return self.conversations[conversation_id]
+            print("✅ Agent and core components initialized successfully.")
+        except Exception as e:
+            print(f"❌ CRITICAL ERROR during Support Agent initialization: {e}")
+            traceback.print_exc()
 
     def answer(self, payload: dict, conversation_id: str) -> dict:
         question = payload.get("question", "")
         live_data_context = payload.get("live_data", "")
-        user_role = payload.get("role", "user")
         
-        memory = self._get_or_create_memory(conversation_id)
+        # Get or create chat history for this conversation
+        chat_history = self.conversations.get(conversation_id, [])
 
         try:
-            # FINAL, POLISHED PROMPT
-            human_friendly_template = """You are Sparky, a helpful AI assistant for Reachify.
+            # This is the modern LangChain Expression Language (LCEL) chain
+            template = """You are Sparky, a helpful AI assistant for Reachify.
 Your job is to provide a direct and concise answer to the user's question.
-Use the Live Data and Context to find the answer. Do not talk about yourself.
+Use the Live Data and Context provided to find the answer. Do not talk about yourself. If the information isn't in the context, say you don't know.
 
-**Live Data (Facts from the user's account):**
+Live Data (Facts from user's account):
 {live_data}
 
-**Context (General Knowledge):**
+Context (General Knowledge from documents):
 {context}
 
-**Chat History:**
+Previous Conversation:
 {chat_history}
 
-**User's Question:** {question}
+User's Question: {question}
 
-**Direct Answer:**
+Direct Answer:
 """
-            
-            final_prompt = PromptTemplate.from_template(human_friendly_template)
-            
-            # ✅ SYNTAX FIX: Removed the space between 'as_' and 'retriever()'
-            retriever = self.vector_store.as_retriever()
-
-            qa_chain = ConversationalRetrievalChain.from_llm(
-                llm=self.langchain_llm_wrapper, 
-                retriever=retriever, 
-                memory=memory, 
-                combine_docs_chain_kwargs={"prompt": final_prompt}
+            prompt = ChatPromptTemplate.from_template(template)
+
+            # Manually format the chat history into a readable string
+            formatted_history = "\n".join([f"Human: {q}\nAssistant: {a}" for q, a in chat_history])
+
+            # The LCEL "pipe"
+            rag_chain = (
+                {
+                    "context": self.retriever | _format_docs_for_context, 
+                    "question": RunnablePassthrough(),
+                    "live_data": lambda x: live_data_context, # Pass live data through
+                    "chat_history": lambda x: formatted_history, # Pass history through
+                }
+                | prompt
+                | self.langchain_llm_wrapper
+                | StrOutputParser()
             )
-            
-            result = qa_chain.invoke({ 
-                "question": question,
-                "live_data": live_data_context
-            })
 
-            raw_answer = result.get("answer", "I'm sorry, I could not find an answer.").strip()
-            
-            final_answer = raw_answer.split("Answer:")[0].split("Direct Answer:")[0].strip()
+            print(f"   - Invoking RAG chain for question: '{question}'")
+            # Invoke the chain by passing just the question string
+            final_answer = rag_chain.invoke(question)
 
-            return {"response": final_answer, "context": format_docs(result.get('source_documents', []))}
+            # Update the conversation memory after getting a successful answer
+            self.conversations[conversation_id] = chat_history + [(question, final_answer)]
+
+            # Get the documents that were used, for transparency
+            source_docs = self.retriever.get_relevant_documents(question)
+            
+            return {
+                "response": final_answer, 
+                "context": _format_docs_for_context(source_docs)
+            }
 
         except Exception as e:
             traceback.print_exc()
-            return {"response": "A critical server error occurred in the AI agent.", "context": str(e)}
-
+            return {
+                "response": "A critical server error occurred in the AI agent.",
+                "context": str(e)
+            }
 
     def generate_caption_variant(self, caption: str, action: str) -> str:
         # Note: You were calling self.llm here but it's defined as self.langchain_llm_wrapper