Update app.py

app.py

@@ -16,6 +16,8 @@ import io
 import docx
 from pptx import Presentation
 import PyPDF2
+import boto3
+from botocore.exceptions import ClientError
 
 # Disable Docker globally
 os.environ["AUTOGEN_USE_DOCKER"] = "0"
@@ -23,24 +25,28 @@ os.environ["AUTOGEN_USE_DOCKER"] = "0"
 class SupplyChainOptimizer:
     def __init__(self):
         # Initialize API keys from environment variables
-        self.openai_api_key = os.environ.get("OPENAI_API_KEY") or os.environ.get("OPENAI_KEY")
+        self.aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
+        self.aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
         self.tavily_api_key = os.environ.get("TAVILY_API_KEY") or os.environ.get("TAVILY_KEY")
         
         # For development/testing, allow demo mode
         self.demo_mode = False
-        if not self.openai_api_key or not self.tavily_api_key:
+        if not self.aws_access_key or not self.aws_secret_key or not self.tavily_api_key:
             print("API keys not found. Running in demo mode.")
             self.demo_mode = True
-            self.openai_api_key = "demo-key"
-            self.tavily_api_key = "demo-key"
-        
-        # Initialize Tavily client (only if not in demo mode)
-        if not self.demo_mode:
+        else:
+            # Initialize Bedrock client
             try:
-                self.tavily = TavilyClient(api_key=self.tavily_api_key)
+                self.bedrock_client = boto3.client(
+                    'bedrock-runtime',
+                    aws_access_key_id=self.aws_access_key,
+                    aws_secret_access_key=self.aws_secret_key,
+                    region_name='us-east-1'  # or your preferred region
+                )
             except Exception as e:
-                print(f"Error initializing Tavily: {e}")
+                print(f"Error initializing Bedrock: {e}")
                 self.demo_mode = True
+
         
         # Initialize agents
         self._setup_agents()
@@ -50,44 +56,74 @@ class SupplyChainOptimizer:
         self.latest_optimization = ""
         self.search_results = ""
 
-    def _setup_agents(self):
-        """Setup AutoGen agents"""
-        self.user_proxy = UserProxyAgent(
-            name="UserProxy",
-            system_message="You are the user interacting with the agents.",
-            human_input_mode="NEVER",
-            code_execution_config={"work_dir": "code", "use_docker": False},
+     def call_claude_api(self, prompt, system_message=""):
+    """Call Claude via AWS Bedrock"""
+    if self.demo_mode:
+        return "Demo mode response"
+    
+    try:
+        body = {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": 4000,
+            "system": system_message,
+            "messages": [
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ]
+        }
+        
+        response = self.bedrock_client.invoke_model(
+            modelId="anthropic.claude-3-haiku-20240307-v1:0",
+            body=json.dumps(body)
         )
+        
+        response_body = json.loads(response['body'].read())
+        return response_body['content'][0]['text']
+        
+    except Exception as e:
+        return f"Error calling Claude API: {str(e)}"   
 
-        # Only setup LLM agents if not in demo mode
-        if not self.demo_mode:
-            self.reasoning_agent = AssistantAgent(
-                name="ReasoningAgent",
-                llm_config={"config_list": [{"model": "gpt-3.5-turbo", "api_key": self.openai_api_key}]},
-                system_message="""You are an expert supply chain analyst. Analyze forecast data using real-time search results.
-                Evaluate if forecasts are reasonable based on current market conditions, events, and trends.
-                Provide detailed reasoning and recommendations. Format your response clearly with proper structure and avoid using asterisks for emphasis."""
-            )
-
-            self.optimization_agent = AssistantAgent(
-                name="OptimizationAgent",
-                llm_config={"config_list": [{"model": "gpt-3.5-turbo", "api_key": self.openai_api_key}]},
-                system_message="""You are a supply chain optimization expert. Create detailed redistribution plans.
-                Consider costs, travel time, inventory levels, and demand forecasts.
-                Provide step-by-step optimization plans with clear recommendations and cost analysis. Use professional formatting without asterisks."""
-            )
+    def _setup_agents(self):
+    """Setup agents with AWS Bedrock Claude"""
+    # Keep user proxy for compatibility but won't be used with direct API calls
+    self.user_proxy = None
+    
+    # Only setup if not in demo mode
+    if not self.demo_mode:
+        # Initialize Bedrock client if not already done
+        if not hasattr(self, 'bedrock_client'):
+            try:
+                self.bedrock_client = boto3.client(
+                    'bedrock-runtime',
+                    aws_access_key_id=self.aws_access_key,
+                    aws_secret_access_key=self.aws_secret_key,
+                    region_name='us-east-1'  # or your preferred region
+                )
+            except Exception as e:
+                print(f"Error initializing Bedrock: {e}")
+                self.demo_mode = True
+                self.reasoning_agent = None
+                self.optimization_agent = None
+                return
+        
+        # Store system messages for direct API calls
+        self.reasoning_system_message = """You are an expert supply chain analyst. Analyze forecast data using real-time search results.
+        Evaluate if forecasts are reasonable based on current market conditions, events, and trends.
+        Provide detailed reasoning and recommendations. Format your response clearly with proper structure and avoid using asterisks for emphasis."""
+        
+        self.optimization_system_message = """You are a supply chain optimization expert. Create detailed redistribution plans.
+        Consider costs, travel time, inventory levels, and demand forecasts.
+        Provide step-by-step optimization plans with clear recommendations and cost analysis. Use professional formatting without asterisks."""
+        
+        # Set agents as enabled (we'll use direct API calls)
+        self.reasoning_agent = "claude-enabled"
+        self.optimization_agent = "claude-enabled"
+    else:
+        self.reasoning_agent = None
+        self.optimization_agent = None
 
-            # Register search tool
-            register_function(
-                self.tavily_search_tool,
-                caller=self.reasoning_agent,
-                executor=self.user_proxy,
-                name="tavily_search_tool",
-                description="Conducts real-time market research using Tavily."
-            )
-        else:
-            self.reasoning_agent = None
-            self.optimization_agent = None
 
     def tavily_search_tool(self, query: Annotated[str, "Market Research Query"]) -> Annotated[str, "Search results"]:
         """Search tool using Tavily API"""
@@ -418,18 +454,9 @@ Note: Set your API keys for AI-powered optimization with real market data
             """
             
             try:
-                with Cache.disk(cache_seed=42) as cache:
-                    reasoning_result = self.user_proxy.initiate_chat(
-                        self.reasoning_agent,
-                        message=reasoning_prompt,
-                        cache=cache,
-                        max_turns=1
-                    )
-                    
-                    if reasoning_result and reasoning_result.chat_history:
-                        self.latest_analysis = reasoning_result.chat_history[-1]['content']
-                    else:
-                        self.latest_analysis = "Analysis completed but no detailed response received."
+                system_msg = "You are an expert supply chain analyst. Analyze forecast data using real-time search results. Evaluate if forecasts are reasonable based on current market conditions, events, and trends. Provide detailed reasoning and recommendations. Format your response clearly with proper structure."
+                
+                self.latest_analysis = self.call_claude_api(reasoning_prompt, system_msg)
             
             except Exception as e:
                 self.latest_analysis = f"Analysis error: {str(e)}"
@@ -457,19 +484,9 @@ Note: Set your API keys for AI-powered optimization with real market data
             """
             
             try:
-                with Cache.disk(cache_seed=43) as cache:
-                    optimization_result = self.user_proxy.initiate_chat(
-                        self.optimization_agent,
-                        message=optimization_prompt,
-                        cache=cache,
-                        max_turns=1
-                    )
-                    
-                    if optimization_result and optimization_result.chat_history:
-                        self.latest_optimization = optimization_result.chat_history[-1]['content']
-                    else:
-                        self.latest_optimization = "Optimization completed but no detailed response received."
-            
+                system_msg = "You are a supply chain optimization expert. Create detailed redistribution plans. Consider costs, travel time, inventory levels, and demand forecasts. Provide step-by-step optimization plans with clear recommendations and cost analysis."
+                
+                self.latest_optimization = self.call_claude_api(optimization_prompt, system_msg)
             except Exception as e:
                 self.latest_optimization = f"Optimization error: {str(e)}"