Upload 2 files

app.py

@@ -6,6 +6,7 @@ from atomic_agents.lib.components.system_prompt_generator import SystemPromptGen
 from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig, BaseAgentInputSchema, BaseAgentOutputSchema
 from dotenv import load_dotenv
 import asyncio
+import httpx
 
 # Load environment variables
 load_dotenv()
@@ -14,6 +15,15 @@ load_dotenv()
 st.title("Math Reasoning Chatbot")
 st.write("Select a provider and chat with the bot to solve math problems!")
 
+# Function to check if the Ollama server is running
+async def check_ollama_health():
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get("http://localhost:11434/v1")
+            return response.status_code == 200
+    except httpx.RequestError:
+        return False
+
 # Function to set up the client based on the chosen provider
 def setup_client(provider):
     if provider == "openai":
@@ -27,23 +37,15 @@ def setup_client(provider):
         display_model = "OpenAI (gpt-4o-mini)"
     elif provider == "ollama":
         from openai import AsyncOpenAI as OllamaClient
+        # Check if Ollama server is running
+        if not asyncio.run(check_ollama_health()):
+            st.error("Ollama server is not running or not accessible at http://localhost:11434. Please try again later or select a different provider.")
+            return None, None, None
         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
@@ -119,15 +121,18 @@ if user_input:
         response_container = st.empty()
         async def stream_response():
             current_response = ""
-            async for partial_response in st.session_state.agent.run_async(input_schema):
-                if hasattr(partial_response, "chat_message") and partial_response.chat_message:
-                    if partial_response.chat_message != current_response:
-                        current_response = partial_response.chat_message
-                        response_container.markdown(current_response)
-
-            # 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", BaseAgentOutputSchema(chat_message=current_response))
+            try:
+                async for partial_response in st.session_state.agent.run_async(input_schema):
+                    if hasattr(partial_response, "chat_message") and partial_response.chat_message:
+                        if partial_response.chat_message != current_response:
+                            current_response = partial_response.chat_message
+                            response_container.markdown(current_response)
+
+                # 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", BaseAgentOutputSchema(chat_message=current_response))
+            except instructor.exceptions.InstructorRetryException as e:
+                response_container.error(f"Failed to connect to the model: {str(e)}. Please try again or select a different provider.")
 
         # Run the async function
         asyncio.run(stream_response())