# 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()