Update agent.py

agent.py

@@ -1,214 +1,440 @@
-"""LangGraph Agent"""
-import os
-from dotenv import load_dotenv
-from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import tools_condition
-from langgraph.prebuilt import ToolNode
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_groq import ChatGroq
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.document_loaders import ArxivLoader
-from langchain_community.vectorstores import SupabaseVectorStore
-from langchain_core.messages import SystemMessage, HumanMessage
-from langchain_core.tools import tool
-from langchain.tools.retriever import create_retriever_tool
-from supabase.client import Client, create_client
+"""
+Modified agent.py - Fixed with Hugging Face models instead of OpenAI
+Fixes LangSmith authentication and missing PostgreSQL function issues
+"""
 
-load_dotenv()
+import os
+import logging
+import warnings
+from typing import List, Dict, Any, Optional, Union
+import pandas as pd
+from supabase import create_client, Client
 
-@tool
-def multiply(a: int, b: int) -> int:
-    """Multiply two numbers.
+# Suppress LangSmith warnings to avoid authentication errors
+warnings.filterwarnings("ignore", category=UserWarning, module="langsmith")
+logging.getLogger("langsmith").setLevel(logging.ERROR)
 
-    Args:
-        a: first int
-        b: second int
-    """
-    return a * b
+# Disable LangSmith tracing to avoid 401 errors
+os.environ["LANGCHAIN_TRACING_V2"] = "false"
 
-@tool
-def add(a: int, b: int) -> int:
-    """Add two numbers.
+try:
+    from langchain.agents import AgentType, AgentExecutor, create_react_agent
+    from langchain.tools import BaseTool, tool
+    from langchain.memory import ConversationBufferMemory
+    from langchain_community.llms import HuggingFacePipeline
+    from langchain_community.embeddings import HuggingFaceEmbeddings
+    from langchain_community.vectorstores import SupabaseVectorStore
+    from langchain.schema import Document
+    from langchain.text_splitter import RecursiveCharacterTextSplitter
+    from langchain.chains import RetrievalQA
+    from langchain.prompts import PromptTemplate
+    from langchain_core.prompts import ChatPromptTemplate
     
-    Args:
-        a: first int
-        b: second int
-    """
-    return a + b
-
-@tool
-def subtract(a: int, b: int) -> int:
-    """Subtract two numbers.
+    # Hugging Face specific imports
+    from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+    from sentence_transformers import SentenceTransformer
+    import torch
     
-    Args:
-        a: first int
-        b: second int
-    """
-    return a - b
+except ImportError as e:
+    print(f"Import error: {e}")
+    print("Please install required packages: pip install transformers sentence-transformers torch")
 
-@tool
-def divide(a: int, b: int) -> int:
-    """Divide two numbers.
-    
-    Args:
-        a: first int
-        b: second int
+class RobotPaiAgent:
     """
-    if b == 0:
-        raise ValueError("Cannot divide by zero.")
-    return a / b
-
-@tool
-def modulus(a: int, b: int) -> int:
-    """Get the modulus of two numbers.
-    
-    Args:
-        a: first int
-        b: second int
+    RobotPai Agent using Hugging Face models instead of OpenAI
+    Fixes authentication and database function issues
     """
-    return a % b
-
-@tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
     
-    Args:
-        query: The search query."""
-    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"wiki_results": formatted_search_docs}
-
-@tool
-def web_search(query: str) -> str:
-    """Search Tavily for a query and return maximum 3 results.
+    def __init__(self, model_name: str = "microsoft/DialoGPT-medium"):
+        print("🤖 Initializing RobotPai Agent with Hugging Face models...")
+        self.model_name = model_name
+        self.setup_environment()
+        self.setup_supabase()
+        self.setup_models()
+        self.setup_vectorstore()
+        self.setup_tools()
+        self.setup_agent()
     
-    Args:
-        query: The search query."""
-    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"web_results": formatted_search_docs}
-
-@tool
-def arvix_search(query: str) -> str:
-    """Search Arxiv for a query and return maximum 3 result.
+    def setup_environment(self):
+        """Setup environment variables with error handling"""
+        # Disable LangSmith to avoid authentication errors
+        os.environ["LANGCHAIN_TRACING_V2"] = "false"
+        
+        # Required environment variables
+        self.supabase_url = os.getenv("SUPABASE_URL")
+        self.supabase_key = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
+        
+        if not all([self.supabase_url, self.supabase_key]):
+            raise ValueError("Missing required environment variables: SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY")
+        
+        print("✅ Environment configured")
     
-    Args:
-        query: The search query."""
-    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"arvix_results": formatted_search_docs}
-
-
-
-# load the system prompt from the file
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    system_prompt = f.read()
+    def setup_supabase(self):
+        """Setup Supabase client and ensure database setup"""
+        try:
+            self.supabase_client: Client = create_client(self.supabase_url, self.supabase_key)
+            self.ensure_database_setup()
+            print("✅ Supabase client initialized")
+        except Exception as e:
+            print(f"⚠️ Supabase setup failed: {e}")
+            self.supabase_client = None
+    
+    def ensure_database_setup(self):
+        """Ensure the database has required tables and functions"""
+        try:
+            # Check if documents table exists
+            result = self.supabase_client.table('documents').select('id').limit(1).execute()
+            print("✅ Documents table exists")
+        except Exception as e:
+            print(f"⚠️ Database setup needed: {e}")
+            print("Please run the SQL setup in your Supabase dashboard:")
+            print("""
+-- Enable pgvector extension
+CREATE EXTENSION IF NOT EXISTS vector;
 
-# System message
-sys_msg = SystemMessage(content=system_prompt)
+-- Create documents table
+CREATE TABLE IF NOT EXISTS documents (
+    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    content TEXT NOT NULL,
+    metadata JSONB DEFAULT '{}',
+    embedding VECTOR(384)  -- Dimension for sentence-transformers
+);
 
-# build a retriever
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") #  dim=768
-supabase: Client = create_client(
-    os.environ.get("SUPABASE_URL"), 
-    os.environ.get("SUPABASE_SERVICE_KEY"))
-vector_store = SupabaseVectorStore(
-    client=supabase,
-    embedding= embeddings,
-    table_name="documents",
-    query_name="match_documents_langchain",
+-- Create match_documents_langchain function
+CREATE OR REPLACE FUNCTION match_documents_langchain(
+    query_embedding VECTOR(384),
+    match_count INT DEFAULT 10,
+    filter JSONB DEFAULT '{}'
 )
-create_retriever_tool = create_retriever_tool(
-    retriever=vector_store.as_retriever(),
-    name="Question Search",
-    description="A tool to retrieve similar questions from a vector store.",
+RETURNS TABLE (
+    id UUID,
+    content TEXT,
+    metadata JSONB,
+    similarity FLOAT
 )
+LANGUAGE plpgsql
+AS $$
+BEGIN
+    RETURN QUERY
+    SELECT
+        documents.id,
+        documents.content,
+        documents.metadata,
+        1 - (documents.embedding <=> query_embedding) AS similarity
+    FROM documents
+    WHERE documents.metadata @> filter
+    ORDER BY documents.embedding <=> query_embedding
+    LIMIT match_count;
+END;
+$$;
+            """)
+    
+    def setup_models(self):
+        """Setup Hugging Face models for LLM and embeddings"""
+        try:
+            # Setup embeddings using sentence-transformers (faster and smaller)
+            print("🔄 Loading embedding model...")
+            self.embeddings = HuggingFaceEmbeddings(
+                model_name="all-MiniLM-L6-v2",  # 384 dimensions, fast and good quality
+                model_kwargs={'device': 'cpu'},  # Use CPU for compatibility
+                encode_kwargs={'normalize_embeddings': True}
+            )
+            print("✅ Embeddings model loaded")
+            
+            # Setup LLM using a lightweight model suitable for HF Spaces
+            print("🔄 Loading language model...")
+            
+            # Use a smaller, faster model for Hugging Face Spaces
+            model_id = "microsoft/DialoGPT-small"  # Smaller model for faster inference
+            
+            try:
+                # Create a text generation pipeline
+                self.llm_pipeline = pipeline(
+                    "text-generation",
+                    model=model_id,
+                    tokenizer=model_id,
+                    max_length=512,
+                    temperature=0.7,
+                    do_sample=True,
+                    device_map="auto" if torch.cuda.is_available() else None,
+                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+                )
+                
+                # Wrap in LangChain HuggingFacePipeline
+                self.llm = HuggingFacePipeline(
+                    pipeline=self.llm_pipeline,
+                    model_kwargs={"temperature": 0.7, "max_length": 512}
+                )
+                print(f"✅ Language model loaded: {model_id}")
+                
+            except Exception as e:
+                print(f"⚠️ Failed to load {model_id}: {e}")
+                # Fallback to a simple text completion
+                print("🔄 Using fallback model...")
+                self.llm = self._create_fallback_llm()
+                
+        except Exception as e:
+            print(f"❌ Model setup failed: {e}")
+            # Create minimal fallback
+            self.embeddings = None
+            self.llm = self._create_fallback_llm()
+    
+    def _create_fallback_llm(self):
+        """Create a simple fallback LLM for when models fail to load"""
+        class SimpleLLM:
+            def __call__(self, prompt: str) -> str:
+                return f"I'm a simple AI assistant. You asked: {prompt[:100]}... I would help you search documents and analyze data, but I need proper model setup."
+            
+            def invoke(self, prompt: str) -> str:
+                return self.__call__(prompt)
+        
+        return SimpleLLM()
+    
+    def setup_vectorstore(self):
+        """Setup vector store with proper error handling"""
+        if not self.supabase_client or not self.embeddings:
+            print("⚠️ Skipping vector store setup - missing dependencies")
+            self.vectorstore = None
+            return
+        
+        try:
+            # Initialize vector store with correct function name
+            self.vectorstore = SupabaseVectorStore(
+                client=self.supabase_client,
+                embedding=self.embeddings,
+                table_name="documents",
+                query_name="match_documents_langchain"  # Use the function we created
+            )
+            print("✅ Vector store initialized")
+            
+        except Exception as e:
+            print(f"⚠️ Vector store setup failed: {e}")
+            self.vectorstore = None
+    
+    def setup_tools(self):
+        """Setup tools for the agent"""
+        self.tools = []
+        
+        # Document Search Tool
+        @tool
+        def search_documents(query: str) -> str:
+            """Search for relevant documents in the knowledge base."""
+            if not self.vectorstore:
+                return "Vector store not available. Please check database setup."
+            
+            try:
+                docs = self.vectorstore.similarity_search(query, k=3)
+                if docs:
+                    results = []
+                    for i, doc in enumerate(docs, 1):
+                        content = doc.page_content[:300] + "..." if len(doc.page_content) > 300 else doc.page_content
+                        results.append(f"Document {i}: {content}")
+                    return "\n\n".join(results)
+                else:
+                    return "No relevant documents found."
+            except Exception as e:
+                return f"Error searching documents: {str(e)}"
+        
+        # CSV Analysis Tool
+        @tool
+        def analyze_csv_data(query: str) -> str:
+            """Analyze CSV data and answer questions about it."""
+            try:
+                # Load the CSV file if it exists
+                if os.path.exists("supabase_docs.csv"):
+                    df = pd.read_csv("supabase_docs.csv")
+                    
+                    # Basic analysis based on query
+                    if "rows" in query.lower() or "count" in query.lower():
+                        return f"The CSV has {len(df)} rows and {len(df.columns)} columns."
+                    elif "columns" in query.lower():
+                        return f"Columns: {', '.join(df.columns.tolist())}"
+                    elif "head" in query.lower() or "first" in query.lower():
+                        return f"First 5 rows:\n{df.head().to_string()}"
+                    else:
+                        return f"CSV loaded with {len(df)} rows. Available columns: {', '.join(df.columns.tolist())}"
+                else:
+                    return "CSV file not found. Please upload supabase_docs.csv"
+            except Exception as e:
+                return f"Error analyzing CSV: {str(e)}"
+        
+        # General Q&A Tool
+        @tool
+        def answer_question(question: str) -> str:
+            """Answer general questions using the language model."""
+            try:
+                # Simple prompt for the question
+                prompt = f"Question: {question}\nAnswer:"
+                response = self.llm.invoke(prompt)
+                return response if isinstance(response, str) else str(response)
+            except Exception as e:
+                return f"I'm unable to process that question right now. Error: {str(e)}"
+        
+        self.tools = [search_documents, analyze_csv_data, answer_question]
+        print(f"✅ {len(self.tools)} tools initialized")
+    
+    def setup_agent(self):
+        """Setup the agent with React framework"""
+        try:
+            # Create a simple prompt template
+            template = """Answer the following questions as best you can. You have access to the following tools:
 
+{tools}
 
+Use the following format:
 
-tools = [
-    multiply,
-    add,
-    subtract,
-    divide,
-    modulus,
-    wiki_search,
-    web_search,
-    arvix_search,
-]
+Question: the input question you must answer
+Thought: you should always think about what to do
+Action: the action to take, should be one of [{tool_names}]
+Action Input: the input to the action
+Observation: the result of the action
+... (this Thought/Action/Action Input/Observation can repeat N times)
+Thought: I now know the final answer
+Final Answer: the final answer to the original input question
 
-# Build graph function
-def build_graph(provider: str = "groq"):
-    """Build the graph"""
-    # Load environment variables from .env file
-    if provider == "google":
-        # Google Gemini
-        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    elif provider == "groq":
-        # Groq https://console.groq.com/docs/models
-        llm = ChatGroq(model="qwen-qwq-32b", temperature=0) # optional : qwen-qwq-32b gemma2-9b-it
-    elif provider == "huggingface":
-        # TODO: Add huggingface endpoint
-        llm = ChatHuggingFace(
-            llm=HuggingFaceEndpoint(
-                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
-                temperature=0,
-            ),
-        )
-    else:
-        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
-    # Bind tools to LLM
-    llm_with_tools = llm.bind_tools(tools)
+Begin!
 
-    # Node
-    def assistant(state: MessagesState):
-        """Assistant node"""
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
-    def retriever(state: MessagesState):
-        """Retriever node"""
-        similar_question = vector_store.similarity_search(state["messages"][0].content)
-        example_msg = HumanMessage(
-            content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
-        )
-        return {"messages": [sys_msg] + state["messages"] + [example_msg]}
+Question: {input}
+Thought: {agent_scratchpad}"""
 
-    builder = StateGraph(MessagesState)
-    builder.add_node("retriever", retriever)
-    builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
-    builder.add_edge(START, "retriever")
-    builder.add_edge("retriever", "assistant")
-    builder.add_conditional_edges(
-        "assistant",
-        tools_condition,
-    )
-    builder.add_edge("tools", "assistant")
+            prompt = PromptTemplate.from_template(template)
+            
+            # Create a simple agent using React pattern
+            if hasattr(self.llm, 'invoke'):
+                agent = create_react_agent(self.llm, self.tools, prompt)
+                self.agent_executor = AgentExecutor(
+                    agent=agent,
+                    tools=self.tools,
+                    verbose=True,
+                    max_iterations=3,
+                    handle_parsing_errors=True,
+                    return_intermediate_steps=True
+                )
+            else:
+                # Fallback for simple LLM
+                self.agent_executor = self._create_simple_executor()
+            
+            print("✅ Agent initialized successfully")
+            
+        except Exception as e:
+            print(f"⚠️ Agent setup failed: {e}")
+            self.agent_executor = self._create_simple_executor()
+    
+    def _create_simple_executor(self):
+        """Create a simple executor when full agent setup fails"""
+        class SimpleExecutor:
+            def __init__(self, tools, llm):
+                self.tools = {tool.name: tool for tool in tools}
+                self.llm = llm
+            
+            def invoke(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+                query = inputs.get("input", "")
+                
+                # Simple routing logic
+                if "document" in query.lower() or "search" in query.lower():
+                    if "search_documents" in self.tools:
+                        result = self.tools["search_documents"].invoke(query)
+                        return {"output": result}
+                
+                elif "csv" in query.lower() or "data" in query.lower():
+                    if "analyze_csv_data" in self.tools:
+                        result = self.tools["analyze_csv_data"].invoke(query)
+                        return {"output": result}
+                
+                else:
+                    if "answer_question" in self.tools:
+                        result = self.tools["answer_question"].invoke(query)
+                        return {"output": result}
+                
+                return {"output": f"I can help you with document search, CSV analysis, or general questions. You asked: {query}"}
+        
+        return SimpleExecutor(self.tools, self.llm)
+    
+    def add_documents(self, texts: List[str], metadatas: List[Dict] = None):
+        """Add documents to the vector store"""
+        if not self.vectorstore:
+            print("⚠️ Vector store not available")
+            return False
+        
+        try:
+            # Split long texts into chunks
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=500,  # Smaller chunks for better performance
+                chunk_overlap=100
+            )
+            
+            all_texts = []
+            all_metadatas = []
+            
+            for i, text in enumerate(texts):
+                chunks = text_splitter.split_text(text)
+                all_texts.extend(chunks)
+                
+                # Add metadata for each chunk
+                base_metadata = metadatas[i] if metadatas and i < len(metadatas) else {}
+                for j, chunk in enumerate(chunks):
+                    chunk_metadata = base_metadata.copy()
+                    chunk_metadata.update({"chunk_id": j, "source_doc": i})
+                    all_metadatas.append(chunk_metadata)
+            
+            # Add to vector store
+            ids = self.vectorstore.add_texts(all_texts, all_metadatas)
+            print(f"✅ Added {len(ids)} document chunks to vector store")
+            return True
+            
+        except Exception as e:
+            print(f"❌ Error adding documents: {e}")
+            return False
+    
+    def process_query(self, query: str) -> str:
+        """Process a user query through the agent"""
+        try:
+            if self.agent_executor:
+                response = self.agent_executor.invoke({"input": query})
+                return response.get("output", "Sorry, I couldn't process your query.")
+            else:
+                return "Agent not properly initialized. Please check your setup."
+        except Exception as e:
+            return f"Error processing query: {str(e)}"
+    
+    def load_csv_for_analysis(self, file_path: str = "supabase_docs.csv") -> bool:
+        """Load CSV data for analysis"""
+        try:
+            if not os.path.exists(file_path):
+                print(f"⚠️ CSV file not found: {file_path}")
+                return False
+            
+            df = pd.read_csv(file_path)
+            print(f"✅ Loaded CSV with {len(df)} rows and {len(df.columns)} columns")
+            
+            # Optionally add CSV content to vector store for searching
+            if self.vectorstore:
+                documents = []
+                for _, row in df.head(100).iterrows():  # Limit to first 100 rows
+                    content = " | ".join([f"{col}: {val}" for col, val in row.items() if pd.notna(val)])
+                    documents.append(content)
+                
+                metadatas = [{"source": "csv_data", "row_id": i} for i in range(len(documents))]
+                self.add_documents(documents, metadatas)
+                print("✅ CSV data added to vector store for searching")
+            
+            return True
+            
+        except Exception as e:
+            print(f"❌ Error loading CSV: {e}")
+            return False
 
-    # Compile graph
-    return builder.compile()
+# Utility function for direct usage
+def create_agent():
+    """Create and return a RobotPai agent instance"""
+    try:
+        agent = RobotPaiAgent()
+        return agent
+    except Exception as e:
+        print(f"Failed to create agent: {e}")
+        return None
 
-# test
-if __name__ == "__main__":
-    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
-    # Build the graph
-    graph = build_graph(provider="groq")
-    # Run the graph
-    messages = [HumanMessage(content=question)]
-    messages = graph.invoke({"messages": messages})
-    for m in messages["messages"]:
-        m.pretty_print()
+# For backward compatibility
+def get_agent():
+    """Get agent instance - for backward compatibility"""
+    return create_agent()