Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,41 +1,52 @@
-# import phase
-import streamlit as st
-from pydantic import BaseModel, Field
 from typing import List, Literal, Dict, Optional
 from agency_swarm import Agent, Agency, set_openai_key, BaseTool
+from pydantic import Field, BaseModel
+import streamlit as st
 
-class AnalyzeProjectRequirement(BaseTool):
-    project_name: str = Field(..., description = "Name of Project")
-    project_description: str = Field(..., description = "Project Description and Goals")
-    project_type: Literal["Web Application", "Mobile App", "API Development", "Data Analytics", "AI/ML Solution", "Other"] = Field(..., description = "Type of project")
-    budget_range: Literal["$10k-$25k", "$25k-$50k", "$100k+"] = Field(..., description = "Budget Range of the project")
+class AnalyzeProjectRequirements(BaseTool):
+    project_name: str = Field(..., description="Name of the project")
+    project_description: str = Field(..., description="Project description and goals")
+    project_type: Literal["Web Application", "Mobile App", "API Development", 
+                         "Data Analytics", "AI/ML Solution", "Other"] = Field(..., 
+                         description="Type of project")
+    budget_range: Literal["$10k-$25k", "$25k-$50k", "$50k-$100k", "$100k+"] = Field(..., 
+                         description="Budget range for the project")
 
     class ToolConfig:
         name = "analyze_project"
-        description = "Analyze project requirements and feasibility"
+        description = "Analyzes project requirements and feasibility"
         one_call_at_a_time = True
 
     def run(self) -> str:
-        """Analyze project and store results in shared state"""
+        """Analyzes project and stores results in shared state"""
         if self._shared_state.get("project_analysis", None) is not None:
-            raise ValueError('Project Analysis already exists, you can move forward to technical specification.')
-
+            raise ValueError("Project analysis already exists. Please proceed with technical specification.")
+        
         analysis = {
             "name": self.project_name,
             "type": self.project_type,
             "complexity": "high",
             "timeline": "6 months",
             "budget_feasibility": "within range",
-            "requirements": ["Scalable Architecture", "Security", "API Integration"]
+            "requirements": ["Scalable architecture", "Security", "API integration"]
         }
-
+        
         self._shared_state.set("project_analysis", analysis)
-        return "Project analysis completed. Please proceed with technical specification"
+        return "Project analysis completed. Please proceed with technical specification."
 
 class CreateTechnicalSpecification(BaseTool):
-    core_technologies: str = Field(..., description = "Comma-seperated list of main technologies and frameworks")
-    architecture_type: Literal["monolithic", "microservice", "serverless", "hybrid"] = Field(..., description = "Proposed Architecture Type")
-    scalability_requirement: Literal("high", "medium", "low") = Field(..., description = "Scalability Needs")
+    architecture_type: Literal["monolithic", "microservices", "serverless", "hybrid"] = Field(
+        ..., 
+        description="Proposed architecture type"
+    )
+    core_technologies: str = Field(
+        ..., 
+        description="Comma-separated list of main technologies and frameworks"
+    )
+    scalability_requirements: Literal["high", "medium", "low"] = Field(
+        ..., 
+        description="Scalability needs"
+    )
 
     class ToolConfig:
         name = "create_technical_spec"
@@ -43,47 +54,51 @@ class CreateTechnicalSpecification(BaseTool):
         one_call_at_a_time = True
 
     def run(self) -> str:
-        """Create technical specification based on the analysis"""
+        """Creates technical specification based on analysis"""
         project_analysis = self._shared_state.get("project_analysis", None)
         if project_analysis is None:
-            raise ValueError("Please analyse project requirements first using the AnalyzeProjectRequirement tool")
-
-        spec ={
-            "project_name" : project_analysis["name"],
+            raise ValueError("Please analyze project requirements first using AnalyzeProjectRequirements tool.")
+        
+        spec = {
+            "project_name": project_analysis["name"],
             "architecture": self.architecture_type,
-            "technologies": self.core_technologies.split(","), 
-            "scalability": self.scalability_requirement
+            "technologies": self.core_technologies.split(","),
+            "scalability": self.scalability_requirements
         }
+        
         self._shared_state.set("technical_specification", spec)
-        return f"Technical Specification created for {project_analysis['name']}"
+        return f"Technical specification created for {project_analysis['name']}."
 
 def init_session_state() -> None:
-    """Initiate session state variable"""
-    if "messages" is not st.session_state:
-        st.session_state.message = []
-    if "api_key" not in st.session_state:
+    """Initialize session state variables"""
+    if 'messages' not in st.session_state:
+        st.session_state.messages = []
+    if 'api_key' not in st.session_state:
         st.session_state.api_key = None
 
-
 def main() -> None:
-    st.set_page_config(page_title = "dev.ai - Place for innovation", layout = "wide")
+    st.set_page_config(page_title="AI Services Agency", layout="wide")
     init_session_state()
-
-    st.title("🎯 dev.ai services - Innovate, Build and Lighten up the future with AI")
-
+    
+    st.title("🚀 AI Services Agency")
+    
     # API Configuration
     with st.sidebar:
-        st.header("API Configuration")
+        st.header("🔑 API Configuration")
         openai_api_key = st.text_input(
             "OpenAI API Key",
-            type = "password",
-            help = "Enter your OpenAI Key that start with sk--xxxxxxxxxxxxxx"
+            type="password",
+            help="Enter your OpenAI API key to continue"
         )
 
         if openai_api_key:
             st.session_state.api_key = openai_api_key
-            st.success("API Key accepted")
+            st.success("API Key accepted!")
         else:
-            st.warning("⚠️ API Key not found, please enter you key to proceed")
+            st.warning("⚠️ Please enter your OpenAI API Key to proceed")
             st.markdown("[Get your API key here](https://platform.openai.com/api-keys)")
             return
+        
+    # Initialize agents with the provided API key
+    set_openai_key(st.session_state.api_key)
+    api_headers = {"Authorization": f"Bearer {st.session_state.api_key}"}