Upload 3 files

Dockerfile

@@ -0,0 +1,27 @@
+# Use an official Python runtime as a parent image
+FROM python:3.11-slim
+
+# Set working directory
+WORKDIR /app
+
+# Copy the requirements file
+COPY requirements.txt .
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Install Ollama
+RUN apt-get update && apt-get install -y curl
+RUN curl -fsSL https://ollama.com/install.sh | sh
+
+# Pre-download the llama3 model
+RUN ollama pull llama3
+
+# Copy the app code
+COPY app.py .
+
+# Expose the Streamlit port
+EXPOSE 8501
+
+# Start Ollama in the background and then run Streamlit
+CMD ollama serve & streamlit run app.py --server.port 8501 --server.address 0.0.0.0

app.py

@@ -0,0 +1,133 @@
+import os
+import streamlit as st
+import instructor
+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
+from dotenv import load_dotenv
+import asyncio
+
+# Load environment variables
+load_dotenv()
+
+# Initialize Streamlit app
+st.title("Math Reasoning Chatbot")
+st.write("Select a provider and chat with the bot to solve math problems!")
+
+# Function to set up the client based on the chosen provider
+def setup_client(provider):
+    if provider == "openai":
+        from openai import AsyncOpenAI
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            st.error("OPENAI_API_KEY not set in environment variables.")
+            return None, None, None
+        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
+        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
+    return client, model, display_model
+
+# Custom system prompt
+system_prompt_generator = SystemPromptGenerator(
+    background=["You are a math genius."],
+    steps=["Think logically step by step and solve a math problem."],
+    output_instructions=[
+        "Summarise your lengthy thinking processes into experienced problems and solutions with thinking order numbers. Do not speak of all the processes.",
+        "Answer in plain English plus formulas.",
+        "Always respond using the proper JSON schema.",
+        "Always use the available additional information and context to enhance the response.",
+    ],
+)
+
+# Provider selection
+providers_list = ["openai", "ollama"]
+selected_provider = st.selectbox("Choose a provider:", providers_list, key="provider_select")
+
+# Set up client and agent based on the selected provider
+client, model, display_model = setup_client(selected_provider)
+if client is None:
+    st.stop()
+
+# Initialize or update the agent
+st.session_state.display_model = display_model
+if "agent" not in st.session_state or st.session_state.get("current_model") != model:
+    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?")
+        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,
+        model=model,
+        system_prompt_generator=system_prompt_generator,
+        memory=st.session_state.memory,
+        system_role="developer",
+    ))
+    st.session_state.current_model = model  # Track the current model to detect changes
+
+# Display the selected model
+st.markdown(f"**Selected Model:** {st.session_state.display_model}")
+
+# Display the system prompt in an expander
+with st.expander("View System Prompt"):
+    system_prompt = system_prompt_generator.generate_prompt()
+    st.text(system_prompt)
+
+# Display conversation history using st.chat_message
+for role, message in st.session_state.conversation:
+    with st.chat_message(role):
+        st.markdown(message)
+
+# User input using st.chat_input
+user_input = st.chat_input(placeholder="e.g., x^4 + a^4 = 0 find cf")
+
+# Process the input and stream the response
+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)
+
+    # Display user message immediately
+    with st.chat_message("user"):
+        st.markdown(user_input)
+
+    # Stream the response
+    with st.chat_message("assistant"):
+        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))
+
+        # Run the async function
+        asyncio.run(stream_response())

requirements.txt

@@ -0,0 +1,4 @@
+streamlit==1.38.0
+instructor==1.5.0
+atomic-agents==0.2.0
+python-dotenv==1.0.1