use langchain directly

app.py

@@ -15,11 +15,8 @@ from langchain_core.tools import tool
 from langchain.agents import AgentExecutor, create_openai_functions_agent
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.schema import SystemMessage
-
-# LangGraph
-from langgraph.graph import END, StateGraph
-from langgraph.prebuilt import ToolNode
-from langgraph.prebuilt.chat_agent import create_agent_executor
+import logging
+import sys
 
 # LangChain
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
@@ -28,10 +25,22 @@ from langchain_core.output_parsers import JsonOutputParser
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
 
-# Load environment variables from .env file
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(sys.stdout)
+    ]
+)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+logger.info("Loading environment variables...")
 load_dotenv()
 
 # Setup FastAPI app
+logger.info("Initializing FastAPI application...")
 app = FastAPI()
 
 
@@ -45,6 +54,7 @@ class AgentState(TypedDict):
 @tool
 def search_web(query: str) -> str:
     """Search the web for information on a given query."""
+    logger.info(f"Searching web for query: {query}")
     # In a real implementation, this would connect to a search API
     return f"Found information about {query}: This is a simulated web search result for '{query}'."
 
@@ -52,39 +62,67 @@ def search_web(query: str) -> str:
 @tool
 def calculate(expression: str) -> str:
     """Calculate the result of a mathematical expression."""
+    logger.info(f"Calculating expression: {expression}")
     try:
         # CAUTION: eval can be dangerous in production; consider using a safer alternative
         return f"Result: {eval(expression)}"
     except Exception as e:
+        logger.error(f"Error calculating expression: {str(e)}")
         return f"Error calculating: {str(e)}"
 
 
 @tool
 def generate_image_prompt(description: str) -> str:
     """Generate a detailed prompt for image generation based on a description."""
+    logger.info(f"Generating image prompt for: {description}")
     enhanced_prompt = f"An image of {description}, high resolution, detailed lighting, professional quality"
     return f"Generated image prompt: {enhanced_prompt}"
 
 
 # Create tools list
+logger.info("Creating tools list...")
 tools = [search_web, calculate, generate_image_prompt]
 
 # Set up the language model
-model = ChatOpenAI(temperature=0.5)
+logger.info("Initializing ChatOpenAI model...")
+try:
+    model = ChatOpenAI(temperature=0.5)
+    logger.info("ChatOpenAI model initialized successfully")
+except Exception as e:
+    logger.error(f"Failed to initialize ChatOpenAI model: {str(e)}")
+    raise
 
 # Create the prompt template
-prompt = ChatPromptTemplate.from_messages([
-    SystemMessage(content="You are a helpful AI assistant with access to tools. Use them when appropriate."),
-    MessagesPlaceholder(variable_name="chat_history"),
-    ("human", "{input}"),
-    MessagesPlaceholder(variable_name="agent_scratchpad"),
-])
+logger.info("Creating prompt template...")
+try:
+    prompt = ChatPromptTemplate.from_messages([
+        SystemMessage(content="You are a helpful AI assistant with access to tools. Use them when appropriate."),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{input}"),
+        MessagesPlaceholder(variable_name="agent_scratchpad"),
+    ])
+    logger.info("Prompt template created successfully")
+except Exception as e:
+    logger.error(f"Failed to create prompt template: {str(e)}")
+    raise
 
 # Create the agent
-agent = create_openai_functions_agent(model, tools, prompt)
+logger.info("Creating OpenAI functions agent...")
+try:
+    agent = create_openai_functions_agent(model, tools, prompt)
+    logger.info("Agent created successfully")
+except Exception as e:
+    logger.error(f"Failed to create agent: {str(e)}")
+    raise
 
 # Create the agent executor
-agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+logger.info("Creating agent executor...")
+try:
+    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+    logger.info("Agent executor created successfully")
+except Exception as e:
+    logger.error(f"Failed to create agent executor: {str(e)}")
+    raise
 
 
 # Define agent nodes
@@ -174,27 +212,41 @@ agent_executor = build_agent_graph()
 # WebSocket for real-time communication
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
+    logger.info("New WebSocket connection established")
     await websocket.accept()
     try:
         while True:
             data = await websocket.receive_text()
-            # Process the message with the agent
-            response = agent_executor.invoke({"input": data, "chat_history": []})
-            await websocket.send_json({"type": "ai_message", "content": response["output"]})
+            logger.info(f"Received message: {data}")
+            try:
+                # Process the message with the agent
+                response = agent_executor.invoke({"input": data, "chat_history": []})
+                logger.info(f"Agent response: {response['output']}")
+                await websocket.send_json({"type": "ai_message", "content": response["output"]})
+            except Exception as e:
+                logger.error(f"Error processing message: {str(e)}")
+                await websocket.send_json({"type": "error", "content": f"Error processing message: {str(e)}"})
     except Exception as e:
-        print(f"Error in WebSocket: {str(e)}")
+        logger.error(f"WebSocket error: {str(e)}")
         await websocket.close()
 
 
 # Serve the HTML frontend
 @app.get("/")
 async def get():
+    logger.info("Serving index.html")
     return FileResponse("index.html")
 
 
 # Mount static files
+logger.info("Mounting static files...")
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    logger.info("Starting uvicorn server...")
+    try:
+        uvicorn.run(app, host="0.0.0.0", port=8000)
+    except Exception as e:
+        logger.error(f"Failed to start uvicorn server: {str(e)}")
+        raise