main.py

# Imports
import os
import sys
import asyncio
from enum import Enum
from typing import Optional, AsyncGenerator, List, Tuple
import gradio as gr
from dotenv import load_dotenv
from langchain_anthropic import ChatAnthropic
from langchain_openai import ChatOpenAI
from langchain_community.llms import HuggingFaceHub
from langchain_community.tools.tavily_search import TavilySearchResults
from psycopg_pool import AsyncConnectionPool
from psycopg.rows import dict_row
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.language_models import BaseChatModel
import psycopg

# NEW IMPORTS FOR PINECONE AND SEMANTIC SEARCH
from pinecone import Pinecone, ServerlessSpec
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore

# Initialize dotenv to load environment variables
load_dotenv() 

# Create an instance of Pinecone using the new API
pinecone_api_key = os.getenv("PINECONE_API_KEY")
pinecone_env = os.getenv("PINECONE_ENVIRONMENT") 
if pinecone_api_key:
    # Create a Pinecone instance. The environment can be passed via ServerlessSpec later.
    pc = Pinecone(api_key=pinecone_api_key)

class LLMProvider(Enum):
    """Enum for supported LLM providers"""
    ANTHROPIC = "anthropic"
    OPENAI = "openai"
    META = "meta"

def get_llm(provider: LLMProvider) -> BaseChatModel:
    """Factory function to create an LLM instance based on the specified provider."""
    if provider == LLMProvider.ANTHROPIC:
        api_key = os.getenv("ANTHROPIC_API_KEY")
        if not api_key:
            raise ValueError("ANTHROPIC_API_KEY environment variable is required for Anthropic")
        return ChatAnthropic(
            api_key=api_key,
            model="claude-3-haiku-20240307",
        )
    
    elif provider == LLMProvider.OPENAI:
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("OPENAI_API_KEY environment variable is required for OpenAI")
        return ChatOpenAI(
            api_key=api_key,
            model="gpt-4o-mini",
            temperature=0
        )
    
    elif provider == LLMProvider.META:
        huggingface_api_key = os.getenv("HUGGINGFACE_API_KEY")
        if not huggingface_api_key:
            raise ValueError("HUGGINGFACE_API_KEY environment variable is required for Meta models")
        return HuggingFaceHub(
            repo_id="meta-llama/Llama-2-70b-chat-hf",
            huggingfacehub_api_token=huggingface_api_key,
            model_kwargs={"temperature": 0.1, "max_length": 4096}
        )
    
    else:
        raise ValueError(f"Unsupported LLM provider: {provider}")

# Initialize Tavily search tool
tavily = TavilySearchResults(max_results=3)

async def check_and_setup_database(pool) -> str:
    """Check if the required tables exist and create them if they don't."""
    async with pool.connection() as conn:
        async with conn.cursor() as cur:
            try:
                await cur.execute("""
                    SELECT EXISTS (
                        SELECT 1 
                        FROM information_schema.tables 
                        WHERE table_schema = 'public' 
                        AND table_name = 'checkpoints'
                    ) as exists;
                """)
                result = await cur.fetchone()
                table_exists = result['exists']
                
                if not table_exists:
                    return "needs_setup"
                return "exists"
                
            except psycopg.Error as e:
                raise Exception(f"Error checking database tables: {str(e)}")

def process_chunk(chunk: dict) -> str:
    """Processes a chunk from the agent and displays information about tool calls or the agent's answer."""
    output = []
    if "agent" in chunk:
        for message in chunk["agent"]["messages"]:
            if "tool_calls" in message.additional_kwargs:
                tool_calls = message.additional_kwargs["tool_calls"]
                for tool_call in tool_calls:
                    tool_name = tool_call["function"]["name"]
                    tool_arguments = eval(tool_call["function"]["arguments"])
                    tool_query = tool_arguments["query"]
                    output.append(f"🔍 Using {tool_name} to search: {tool_query}")
            else:
                output.append(f"🤖 {message.content}")
    return "\n".join(output)

def ingest_document(file_obj) -> PineconeVectorStore:
    """
    Ingest uploaded document into Pinecone vector store for semantic search.
    Supports .txt, .docx, .doc, and .pdf files with robust fallback mechanisms.
    
    Args:
        file_obj: Can be a string file path, bytes object, or file-like object
        
    Returns:
        PineconeVectorStore: The vector store containing the embedded document chunks
    """
    import os
    import tempfile
    from typing import Optional
    
    def extract_pdf_text(pdf_path) -> str:
        """Extract text from PDF using multiple methods with fallbacks."""
        text_content = ""
        
        # Method 1: PyPDF2
        try:
            import PyPDF2
            with open(pdf_path, 'rb') as pdf_file:
                pdf_reader = PyPDF2.PdfReader(pdf_file)
                for page_num in range(len(pdf_reader.pages)):
                    page = pdf_reader.pages[page_num]
                    text = page.extract_text()
                    if text:
                        text_content += text + "\n"
            if text_content.strip():
                print("Successfully extracted text with PyPDF2")
                return text_content
        except Exception as e:
            print(f"PyPDF2 extraction failed: {e}")
        
        # Method 2: pdfplumber
        try:
            import pdfplumber
            with pdfplumber.open(pdf_path) as pdf:
                pages_text = []
                for page in pdf.pages:
                    text = page.extract_text()
                    if text:
                        pages_text.append(text)
                text_content = "\n".join(pages_text)
            if text_content.strip():
                print("Successfully extracted text with pdfplumber")
                return text_content
        except Exception as e:
            print(f"pdfplumber extraction failed: {e}")
        
        # Method 3: textract
        try:
            import textract
            text_content = textract.process(pdf_path).decode('utf-8', errors='ignore')
            if text_content.strip():
                print("Successfully extracted text with textract")
                return text_content
        except Exception as e:
            print(f"textract extraction failed: {e}")
        
        # Method 4: OCR for scanned PDFs
        try:
            import pytesseract
            from pdf2image import convert_from_path
            
            # Convert PDF to images
            images = convert_from_path(pdf_path)
            ocr_texts = []
            
            # Extract text from each image using OCR
            for i, image in enumerate(images):
                print(f"Processing page {i+1} with OCR")
                ocr_texts.append(pytesseract.image_to_string(image))
            
            text_content = "\n".join(ocr_texts)
            if text_content.strip():
                print("Successfully extracted text using OCR")
                return text_content
        except Exception as e:
            print(f"OCR extraction failed: {e}")
        
        return text_content

    def extract_docx_text(docx_path) -> str:
        """Extract text from .docx files."""
        try:
            import docx
            doc = docx.Document(docx_path)
            return "\n".join([paragraph.text for paragraph in doc.paragraphs])
        except Exception as e:
            print(f"DOCX extraction failed: {e}")
            return ""

    def extract_doc_text(doc_path) -> str:
        """Extract text from .doc files."""
        try:
            import textract
            return textract.process(doc_path).decode('utf-8', errors='ignore')
        except Exception as e:
            print(f"DOC extraction failed: {e}")
            
            # Alternative method for .doc files
            try:
                from subprocess import Popen, PIPE
                # Using antiword if available
                process = Popen(['antiword', doc_path], stdout=PIPE)
                stdout, stderr = process.communicate()
                return stdout.decode('utf-8', errors='ignore')
            except Exception as alt_e:
                print(f"Alternative DOC extraction failed: {alt_e}")
                return ""

    def save_bytes_to_temp_file(bytes_data) -> Optional[str]:
        """Save bytes data to a temporary file and return the path."""
        try:
            with tempfile.NamedTemporaryFile(delete=False) as temp:
                temp.write(bytes_data)
                return temp.name
        except Exception as e:
            print(f"Failed to save bytes to temp file: {e}")
            return None

    try:
        file_path = None
        file_content = None
        temp_file_created = False
        
        # Handle string file path
        if isinstance(file_obj, str):
            file_path = file_obj
        # Handle bytes or file-like object
        elif isinstance(file_obj, bytes):
            file_path = save_bytes_to_temp_file(file_obj)
            temp_file_created = True
        elif hasattr(file_obj, 'read') and callable(file_obj.read):
            # File-like object
            file_content = file_obj.read()
            if isinstance(file_content, bytes):
                file_path = save_bytes_to_temp_file(file_content)
                temp_file_created = True
        
        # Extract content based on file type
        if file_path:
            file_lower = file_path.lower()
            
            # Handle PDF files
            if file_lower.endswith('.pdf'):
                file_content = extract_pdf_text(file_path)
            
            # Handle Word documents (.docx)
            elif file_lower.endswith('.docx'):
                file_content = extract_docx_text(file_path)
            
            # Handle Word documents (.doc)
            elif file_lower.endswith('.doc'):
                file_content = extract_doc_text(file_path)
            
            # Handle plain text files
            elif file_lower.endswith('.txt'):
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    file_content = f.read()
            
            # Handle markdown or other text files
            elif file_lower.endswith('.md'):
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    file_content = f.read()
            
            # Handle unknown file types
            else:
                try:
                    import textract
                    file_content = textract.process(file_path).decode('utf-8', errors='ignore')
                except Exception as e:
                    print(f"Textract extraction for unknown file type failed: {e}")
                    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                        file_content = f.read()
        
        # Clean up temp file if created
        if temp_file_created and file_path:
            try:
                os.unlink(file_path)
            except Exception as e:
                print(f"Failed to remove temp file: {e}")
        
        # Validate extracted content
        if not file_content or file_content.strip() == "":
            file_size = os.path.getsize(file_path) if file_path else "unknown"
            file_type = file_path.split('.')[-1] if file_path else "unknown"
            raise ValueError(f"Could not extract text from the document. File size: {file_size} bytes, type: {file_type}")
        
        # Split content into chunks with some overlap for context preservation
        chunk_size = 500
        overlap = 50
        chunks = []
        
        for i in range(0, len(file_content), chunk_size - overlap):
            chunk = file_content[i:i + chunk_size]
            if len(chunk.strip()) > 0:  # Only add non-empty chunks
                chunks.append(chunk)
        
        # Import required libraries
        from langchain.embeddings.openai import OpenAIEmbeddings
        import pinecone
        from pinecone import Pinecone, ServerlessSpec
        from langchain_pinecone import PineconeVectorStore
        
        # Initialize Pinecone
        pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
        pinecone_env = os.getenv("PINECONE_ENVIRONMENT", "us-west1-gcp")
        
        # Initialize embeddings
        embeddings = OpenAIEmbeddings(api_key=os.getenv("OPENAI_API_KEY"))
        index_name = os.getenv("PINECONE_INDEX_NAME", "document-index")
        
        # Create the Pinecone index if it does not exist
        if index_name not in pc.list_indexes().names():
            pc.create_index(
                name=index_name,
                dimension=1536,  # Dimension for OpenAI embeddings
                metric="cosine",
                spec=ServerlessSpec(cloud='aws', region=pinecone_env)
            )
        
        # Get a reference to the index
        index = pc.Index(index_name)

        # Create and populate vector store
        vectorstore = PineconeVectorStore(
            index=index,
            embedding=embeddings,
            text_key="text"
        )
        
        # Insert documents with metadata
        texts_with_metadata = []
        for i, chunk in enumerate(chunks):
            metadata = {
                "chunk_id": i,
                "source": file_path if file_path else "uploaded_document",
                "chunk_count": len(chunks)
            }
            texts_with_metadata.append((chunk, metadata))
        
        # Add texts with metadata
        vectorstore.add_texts([t[0] for t in texts_with_metadata], 
                             metadatas=[t[1] for t in texts_with_metadata])

        print(f"Successfully ingested document with {len(chunks)} chunks")
        return vectorstore
    
    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        raise Exception(f"Error processing document: {str(e)}\nDetails: {error_details}")

class AgentInterface:
    def __init__(self):
        self.pool = None
        self.conn = None
        self.memory = None
        self.langgraph_agent = None
        self.chat_history = []
        self.vectorstore = None  # For Pinecone semantic search
        
    async def initialize(self, provider: str) -> str:
        """Initialize the agent with the selected provider."""
        try:
            llm_provider = LLMProvider(provider.lower())
            llm = get_llm(llm_provider)
            
            # Initialize database connection pool
            self.pool = AsyncConnectionPool(
                conninfo=os.getenv("POSTGRES_CONNECTION_STRING"),
                max_size=20,
                kwargs={
                    "autocommit": True,
                    "prepare_threshold": 0,
                    "row_factory": dict_row,
                    "keepalives": 1,
                    "keepalives_idle": 30,  # Seconds of inactivity before sending keepalive
                    "keepalives_interval": 10,  # Seconds between keepalives if not acknowledged
                    "keepalives_count": 5  # Number of unacknowledged keepalives before considering connection dead
                }
            )
            
            # Check database setup
            db_status = await check_and_setup_database(self.pool)
            
            # Initialize memory with a new connection from the pool
            async with self.pool.connection() as conn:
                self.memory = AsyncPostgresSaver(conn)
                if db_status == "needs_setup":
                    await self.memory.setup()
                    status_msg = "Database tables created successfully!"
                else:
                    status_msg = "Database tables already exist!"
            
            # Initialize agent with the Tavily search tool
            self.langgraph_agent = create_react_agent(
                model=llm, 
                tools=[tavily], 
                checkpointer=self.memory
            )
            
            return f"✅ Initialized {provider} successfully!\n{status_msg}"
            
        except Exception as e:
            if self.pool:
                await self.pool.close()
            raise Exception(str(e))
    
    async def chat(self, message: str, history: List[Tuple[str, str]]) -> AsyncGenerator[List[Tuple[str, str]], None]:
        """Process a chat message and yield responses."""
        if not self.langgraph_agent:
            yield history + [(message, "❌ Please initialize the agent first by selecting a provider.")]
            return
        
        # If a document has been ingested, retrieve relevant context using semantic search
        if self.vectorstore is not None:
            search_results = self.vectorstore.similarity_search(message)
            if search_results:
                context = "\n".join([doc.page_content for doc in search_results])
                message = f"Context: {context}\n\nQuestion: {message}"
        
        try:
            response_parts = []
            async for chunk in self.langgraph_agent.astream(
                {"messages": [HumanMessage(content=message)]},
                {"configurable": {"thread_id": "1"}},
            ):
                chunk_text = process_chunk(chunk)
                if chunk_text:
                    response_parts.append(chunk_text)
                    yield history + [(message, "\n".join(response_parts))]
            
        except Exception as e:
            yield history + [(message, f"❌ Error processing message: {str(e)}")]
    
    async def cleanup(self):
        """Clean up resources."""
        if self.pool:
            await self.pool.close()

async def initialize_agent(provider: str):
    """Initialize the agent interface."""
    agent = AgentInterface()
    try:
        status = await agent.initialize(provider)
        return agent, status
    except Exception as e:
        await agent.cleanup()
        return None, f"❌ Error initializing agent: {str(e)}"

async def chat_response(message: str, history: List[Tuple[str, str]], agent: AgentInterface) -> AsyncGenerator[List[Tuple[str, str]], None]:
    """Handle chat responses."""
    if not agent or not agent.langgraph_agent:
        yield history + [(message, "Please initialize the agent first.")]
        return
    
    async for updated_history in agent.chat(message, history):
        yield updated_history

def ingest_and_update(agent, file):
    """Ingest document file and update the agent's vectorstore."""
    if not agent:
        return "Agent is not initialized."
    try:
        vectorstore = ingest_document(file)
        agent.vectorstore = vectorstore
        return "Document ingested successfully!"
    except Exception as e:
        error_msg = str(e)
        if "utf-8" in error_msg and "decode" in error_msg:
            return "Error: The file appears to be in a binary format. Please ensure you're uploading a valid document file."
        return f"Error in document ingestion: {error_msg}"

def create_interface():
    """Create and launch the Gradio interface."""
    with gr.Blocks(theme=gr.themes.Ocean()) as demo:
        agent_state = gr.State()
        
        with gr.Row():
            with gr.Column(scale=1):
                status = gr.Textbox(label="Status", interactive=False)
                provider = gr.Dropdown(
                    choices=["Anthropic", "OpenAI", "Meta"],
                    label="Select LLM Provider",
                    value="OpenAI"
                )
                init_btn = gr.Button("Initialize Agent")
            with gr.Column(scale=2):
                ingestion_status = gr.Textbox(label="Ingestion Status", interactive=False)
                file_input = gr.File(label="Upload Document", file_count="single", type="filepath", file_types=[".txt", ".pdf", ".docx", ".doc", ".md"])
                upload_btn = gr.Button("Ingest Document")
            
        
        with gr.Row():
            chatbot = gr.Chatbot(
                label="Chat History",
                height=600,
                show_copy_button=True
            )
        
        with gr.Row():
            with gr.Column(scale=2):
                msg = gr.Textbox(
                    label="Your message",
                    placeholder="Type your message here...",
                    lines=2
                )
            with gr.Column(scale=1):
                submit_btn = gr.Button("Send", variant="primary")
        
        async def init_and_store(provider_name):
            agent, init_status = await initialize_agent(provider_name)
            return agent, init_status
        
        init_btn.click(
            init_and_store,
            inputs=[provider],
            outputs=[agent_state, status]
        )
        
        upload_btn.click(
            fn=ingest_and_update,
            inputs=[agent_state, file_input],
            outputs=ingestion_status
        )
        
        submit_btn.click(
            chat_response,
            inputs=[msg, chatbot, agent_state],
            outputs=chatbot
        ).then(
            lambda: "",
            outputs=msg
        )
        
        msg.submit(
            chat_response,
            inputs=[msg, chatbot, agent_state],
            outputs=chatbot
        ).then(
            lambda: "",
            outputs=msg
        )
    
    return demo

if __name__ == "__main__":
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    
    demo = create_interface()
    demo.queue()
    #demo.launch(pwa=True, share=False)
    demo.launch()