Upload 2 files

app.py

@@ -1,7 +1,6 @@
 import os
 import streamlit as st
 import instructor
-import httpx
 from atomic_agents.lib.components.agent_memory import AgentMemory
 from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator
 from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig, BaseAgentInputSchema, BaseAgentOutputSchema
@@ -23,22 +22,28 @@ def setup_client(provider):
         if not api_key:
             st.error("OPENAI_API_KEY not set in environment variables.")
             return None, None, None
-        # Create a custom HTTP client without proxies
-        http_client = httpx.AsyncClient()
-        # Initialize AsyncOpenAI with the custom HTTP client
-        openai_client = AsyncOpenAI(api_key=api_key, http_client=http_client)
-        client = instructor.from_openai(openai_client)
+        client = instructor.from_openai(AsyncOpenAI(api_key=api_key))
         model = "gpt-4o-mini"
         display_model = "OpenAI (gpt-4o-mini)"
     elif provider == "ollama":
         from openai import AsyncOpenAI as OllamaClient
-        # Create a custom HTTP client without proxies
-        http_client = httpx.AsyncClient()
-        # Initialize OllamaClient with the custom HTTP client
-        ollama_client = OllamaClient(base_url="http://localhost:11434/v1", api_key="ollama", http_client=http_client)
-        client = instructor.from_openai(ollama_client, mode=instructor.Mode.JSON)
+        client = instructor.from_openai(
+            OllamaClient(base_url="http://localhost:11434/v1", api_key="ollama"), mode=instructor.Mode.JSON
+        )
         model = "llama3"
         display_model = "Ollama (llama3)"
+    # elif provider == "gemini":
+    #     from openai import AsyncOpenAI
+    #     api_key = os.getenv("GEMINI_API_KEY")
+    #     if not api_key:
+    #         st.error("GEMINI_API_KEY not set in environment variables.")
+    #         return None, None, None
+    #     client = instructor.from_openai(
+    #         AsyncOpenAI(api_key=api_key, base_url="https://generativelanguage.googleapis.com/v1beta/openai/"),
+    #         mode=instructor.Mode.JSON,
+    #     )
+    #     model = "gemini-2.0-flash-exp"
+    #     display_model = "Gemini (gemini-2.0-flash-exp)"
     else:
         st.error(f"Unsupported provider: {provider}")
         return None, None, None
@@ -71,8 +76,7 @@ if "agent" not in st.session_state or st.session_state.get("current_model") != m
     if "memory" not in st.session_state:
         st.session_state.memory = AgentMemory()
         initial_message = BaseAgentOutputSchema(chat_message="Hello! I'm here to help with math problems. What can I assist you with today?")
-        # Pass the chat_message string instead of the BaseAgentOutputSchema object
-        st.session_state.memory.add_message("assistant", initial_message.chat_message)
+        st.session_state.memory.add_message("assistant", initial_message)
         st.session_state.conversation = [("assistant", initial_message.chat_message)]
     st.session_state.agent = BaseAgent(config=BaseAgentConfig(
         client=client,
@@ -104,7 +108,7 @@ if user_input:
     # Add user message to conversation and memory
     st.session_state.conversation.append(("user", user_input))
     input_schema = BaseAgentInputSchema(chat_message=user_input)
-    st.session_state.memory.add_message("user", input_schema.chat_message)  # Pass the chat_message string
+    st.session_state.memory.add_message("user", input_schema)
 
     # Display user message immediately
     with st.chat_message("user"):
@@ -123,7 +127,7 @@ if user_input:
 
             # After streaming completes, add the final response to conversation and memory
             st.session_state.conversation.append(("assistant", current_response))
-            st.session_state.memory.add_message("assistant", current_response)  # Pass the string directly
+            st.session_state.memory.add_message("assistant", BaseAgentOutputSchema(chat_message=current_response))
 
         # Run the async function
         asyncio.run(stream_response())

requirements.txt

@@ -1,6 +1,6 @@
-streamlit==1.38.0
-instructor==1.3.4
-atomic-agents==0.2.1
+streamlit==1.44.1
+instructor==1.6.3
+-e git+https://github.com/BrainBlend-AI/atomic-agents.git@283da044f6a22b993dae0751c08dbae52ab2d2df#egg=atomic_agents
 python-dotenv==1.0.1
-openai==1.35.12
-httpx==0.27.0
+openai==1.54.3
+httpx==0.27.2