o1-mini

app.py

@@ -9,6 +9,7 @@ import json
 import queue
 import logging
 from PIL import Image
+from typing import Optional
 
 # ------------------------
 # LangSmith imports
@@ -18,12 +19,15 @@ from langsmith.wrappers import wrap_openai
 from langsmith import traceable
 
 # ------------------------
-# Configure logging (optional but recommended)
+# Configure logging
 # ------------------------
 def init_logging():
     logging.basicConfig(
-        format="[%(asctime)s] %(levelname)+8s: %(message)s",
+        format="[%(asctime)s] %(levelname)s: %(message)s",
         level=logging.INFO,
+        handlers=[
+            logging.StreamHandler()
+        ]
     )
     return logging.getLogger()
 
@@ -37,7 +41,9 @@ api_key = os.getenv("OPENAI_API_KEY")
 assistant_id = os.getenv("ASSISTANT_ID_SOLUTION_SPECIFIER_A")  # The assistant we want to call
 
 if not api_key or not assistant_id:
-    raise RuntimeError("Please set OPENAI_API_KEY and ASSISTANT_ID_SOLUTION_SPECIFIER_A in your environment")
+    logger.error("Environment variables OPENAI_API_KEY and ASSISTANT_ID_SOLUTION_SPECIFIER_A must be set.")
+    st.error("Missing environment configuration. Please set the required environment variables.")
+    st.stop()
 
 # ------------------------
 # Wrap the OpenAI client for LangSmith traceability
@@ -46,168 +52,235 @@ openai_client = openai.Client(api_key=api_key)
 client = wrap_openai(openai_client)
 
 # ------------------------
-# Streamlit session state
+# Streamlit session state initialization
 # ------------------------
 if "messages" not in st.session_state:
     st.session_state["messages"] = []
 
-if "thread" not in st.session_state:
-    st.session_state["thread"] = None
+if "thread_id" not in st.session_state:
+    st.session_state["thread_id"] = None
 
 if "tool_requests" not in st.session_state:
     st.session_state["tool_requests"] = queue.Queue()
 
+if "current_run" not in st.session_state:
+    st.session_state["current_run"] = None
+
 tool_requests = st.session_state["tool_requests"]
 
 # ------------------------
-# Utility to remove citations like: 【12†somefile】
-# You can adapt to your own "annotations" handling if needed
+# Utility to remove citations like: 
 # ------------------------
 def remove_citation(text: str) -> str:
     pattern = r"【\d+†\w+】"
     return re.sub(pattern, "📚", text)
 
+# ------------------------
+# Function to handle tool requests (function calls)
+# ------------------------
+def handle_tool_request(event):
+    """
+    Processes function call requests from the assistant.
+    """
+    logger.info(f"Handling tool request: {event}")
+    st.toast("Processing a function call...", icon=":hammer_and_wrench:")
+    tool_outputs = []
+    data = event.data
+
+    for tool_call in data.required_action.submit_tool_outputs.tool_calls:
+        function_name = tool_call.function.name
+        arguments = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
+
+        logger.info(f"Executing function '{function_name}' with arguments {arguments}")
+
+        try:
+            # Map function names to actual implementations
+            if function_name == "hello_world":
+                output = hello_world(**arguments)
+            elif function_name == "another_function":
+                output = another_function(**arguments)
+            else:
+                raise ValueError(f"Unrecognized function name: {function_name}")
+
+            tool_outputs.append({"tool_call_id": tool_call.id, "output": output})
+            logger.info(f"Function '{function_name}' executed successfully.")
+
+        except Exception as e:
+            logger.error(f"Error executing function '{function_name}': {e}")
+            error_response = {"status": "error", "message": str(e)}
+            tool_outputs.append({"tool_call_id": tool_call.id, "output": json.dumps(error_response)})
+
+    st.toast("Function call completed.", icon=":white_check_mark:")
+    return tool_outputs, data.thread_id, data.id
+
+# ------------------------
+# Example function implementations
+# ------------------------
+def hello_world(name: str = "World") -> str:
+    """
+    Example function that returns a greeting.
+    """
+    time.sleep(2)  # Simulate a delay for a long-running task
+    return f"Hello, {name}! This message is from a function call."
+
+def another_function(param1: str, param2: int) -> str:
+    """
+    Another example function.
+    """
+    time.sleep(1)
+    return f"Received param1: {param1} and param2: {param2}."
+
+# ------------------------
+# Streamlit UI Components
+# ------------------------
+def display_message(role: str, content: str):
+    """
+    Displays a message in the Streamlit chat interface.
+    """
+    with st.chat_message(role):
+        if role == "assistant" and isinstance(content, Image.Image):
+            st.image(content)
+        else:
+            st.write(content)
+
 # ------------------------
 # Helper: data streamer for text & images
-# Adapted from the Medium article approach
-# to handle text deltas, images, or function calls
 # ------------------------
 def data_streamer():
     """
     Streams data from the assistant run. Yields text or images
     and enqueues tool requests (function calls) to tool_requests.
     """
-    st.toast("Thinking...", icon=":material/emoji_objects:")
+    logger.info("Starting data streamer.")
+    st.toast("Thinking...", icon=":hourglass_flowing_sand:")
     content_produced = False
-
-    for event in st.session_state["run_stream"]:
-        match event.event:
-            case "thread.message.delta":
-                # A chunk of text or an image
-                content = event.data.delta.content[0]
-                match content.type:
-                    case "text":
-                        text_value = content.text.value
-                        content_produced = True
-                        # Optionally remove citations, etc.
-                        yield remove_citation(text_value)
-
-                    case "image_file":
-                        # If the assistant returns an image
-                        file_id = content.image_file.file_id
-                        content_produced = True
-                        image_content = io.BytesIO(client.files.content(file_id).read())
-                        yield Image.open(image_content)
-
-            case "thread.run.requires_action":
-                # The assistant is requesting a function call
-                logger.info(f"[Tool Request] {event}")
-                tool_requests.put(event)
-                if not content_produced:
-                    # We can yield a placeholder if the model hasn't said anything yet
-                    yield "[LLM is requesting a function call]"
-                return
-
-            case "thread.run.failed":
-                # The run failed for some reason
-                logger.error(f"Run failed: {event}")
-                return
-
-    # If we successfully streamed everything
-    st.toast("Completed", icon=":material/emoji_objects:")
+    accumulated_content = ""
+
+    try:
+        for event in st.session_state["current_run"]:
+            match event.event:
+                case "thread.message.delta":
+                    content = event.data.delta.content[0]
+                    match content.type:
+                        case "text":
+                            text_value = content.text.value
+                            accumulated_content += text_value
+                            content_produced = True
+                            yield remove_citation(text_value)
+
+                        case "image_file":
+                            file_id = content.image_file.file_id
+                            logger.info(f"Received image file ID: {file_id}")
+                            image_content = io.BytesIO(client.files.content(file_id).read())
+                            image = Image.open(image_content)
+                            yield image
+
+                case "thread.run.requires_action":
+                    logger.info(f"Run requires action: {event}")
+                    tool_requests.put(event)
+                    if not content_produced:
+                        yield "[LLM is requesting a function call...]"
+                    return
+
+                case "thread.run.failed":
+                    logger.error(f"Run failed: {event}")
+                    st.error("The assistant encountered an error and couldn't complete the request.")
+                    return
+
+    except Exception as e:
+        logger.exception(f"Exception in data_streamer: {e}")
+        st.error(f"An unexpected error occurred: {e}")
+
+    finally:
+        st.toast("Completed", icon=":checkered_flag:")
 
 # ------------------------
 # Helper: display the streaming content
-# This wraps data_streamer in st.write_stream
-# so you can see partial tokens in real-time
 # ------------------------
 def display_stream(run_stream, create_context=True):
     """
     Grabs tokens from data_streamer() and displays them in real-time.
     If `create_context=True`, messages are displayed as an assistant block.
     """
-    st.session_state["run_stream"] = run_stream
+    logger.info("Displaying stream.")
+    st.session_state["current_run"] = run_stream
+
     if create_context:
         with st.chat_message("assistant"):
-            streamed_result = st.write_stream(data_streamer)
+            for content in data_streamer():
+                display_message("assistant", content)
     else:
-        streamed_result = st.write_stream(data_streamer)
-
-    # Return whatever the final token stream is
-    return streamed_result
+        for content in data_streamer():
+            display_message("assistant", content)
 
-# ------------------------
-# Example of handling a function call (requires_action)
-# If your Assistant uses function calling (e.g. code interpreter),
-# you'd parse arguments, run the function, and return output here.
-# ------------------------
-def handle_tool_request(event):
-    """
-    Demonstrates how you might handle a function call.
-    In practice, you'd parse the arguments from the event
-    and run your custom logic. Then return outputs as JSON.
-    """
-    st.toast("Running a function (this is user-defined code)", icon=":material/function:")
-    tool_outputs = []
-    data = event.data
-    for tool_call in data.required_action.submit_tool_outputs.tool_calls:
-        if tool_call.function.arguments:
-            function_args = json.loads(tool_call.function.arguments)
-        else:
-            function_args = {}
-
-        match tool_call.function.name:
-            case "hello_world":
-                # Example: implement a user-defined function
-                name = function_args.get("name", "anonymous")
-                time.sleep(2)  # Simulate a long-running function
-                output_val = f"Hello, {name}! This was from a local function."
-                tool_outputs.append({"tool_call_id": tool_call.id, "output": output_val})
-            case _:
-                # If unknown function name
-                msg = {"status": "error", "message": "Unknown function request."}
-                tool_outputs.append({"tool_call_id": tool_call.id, "output": json.dumps(msg)})
-    return tool_outputs, data.thread_id, data.id
+    # After streaming, accumulate the final content
+    # This assumes that the entire content has been yielded
+    # You might want to enhance this to handle partial content or interruptions
+    # Here, we simply capture accumulated content if it's text
+    # For images, it's already displayed
+    if accumulated_text := remove_citation(accumulated_content.strip()):
+        st.session_state["messages"].append({"role": "assistant", "content": accumulated_text})
 
 # ------------------------
-# Main chat logic
+# Main chat logic with traceability
 # ------------------------
-@traceable  # Make this function traceable via LangSmith
+@traceable  # Enable LangSmith traceability
 def generate_assistant_reply(user_input: str):
     """
-    1. If no thread exists, create a new one.
-    2. Insert user message into the thread.
-    3. Use the Assistants API to create a run + stream the response.
-    4. If the assistant requests a function call, handle it and stream again.
+    Handles user input by creating or continuing a thread,
+    sending the message to the assistant, and streaming the response.
     """
+    logger.info(f"User input received: {user_input}")
+
     # Create or retrieve thread
-    if not st.session_state["thread"]:
-        st.session_state["thread"] = client.beta.threads.create()
-    thread = st.session_state["thread"]
+    if not st.session_state["thread_id"]:
+        logger.info("Creating a new thread.")
+        thread = client.beta.threads.create()
+        st.session_state["thread_id"] = thread.id
+    else:
+        thread = client.beta.threads.retrieve(thread_id=st.session_state["thread_id"])
+        logger.info(f"Using existing thread ID: {thread.id}")
 
     # Add user message to the thread
-    client.beta.threads.messages.create(
-        thread_id=thread.id,
-        role="user",
-        content=user_input
-    )
-
-    # Start streaming assistant response
-    with client.beta.threads.runs.stream(
-        thread_id=thread.id,
-        assistant_id=assistant_id,
-    ) as run_stream:
-        display_stream(run_stream)
-
-    # If the assistant requested any tool calls, handle them now
+    try:
+        client.beta.threads.messages.create(
+            thread_id=thread.id,
+            role="user",
+            content=user_input
+        )
+        logger.info("User message added to thread.")
+    except Exception as e:
+        logger.exception(f"Failed to add user message to thread: {e}")
+        st.error("Failed to send your message. Please try again.")
+        return
+
+    # Create and stream assistant response
+    try:
+        with client.beta.threads.runs.stream(
+            thread_id=thread.id,
+            assistant_id=assistant_id,
+        ) as run_stream:
+            st.session_state["current_run"] = run_stream
+            display_stream(run_stream)
+    except Exception as e:
+        logger.exception(f"Failed to stream assistant response: {e}")
+        st.error("Failed to receive a response from the assistant. Please try again.")
+
+    # Handle any function calls requested by the assistant
     while not tool_requests.empty():
         event = tool_requests.get()
         tool_outputs, t_id, run_id = handle_tool_request(event)
-        # Submit tool outputs
-        with client.beta.threads.runs.submit_tool_outputs_stream(
-            thread_id=t_id, run_id=run_id, tool_outputs=tool_outputs
-        ) as next_stream:
-            display_stream(next_stream, create_context=False)
+
+        try:
+            with client.beta.threads.runs.submit_tool_outputs_stream(
+                thread_id=t_id,
+                run_id=run_id,
+                tool_outputs=tool_outputs
+            ) as tool_stream:
+                display_stream(tool_stream, create_context=False)
+        except Exception as e:
+            logger.exception(f"Failed to submit tool outputs: {e}")
+            st.error("Failed to process a function call from the assistant.")
 
 # ------------------------
 # Streamlit UI
@@ -218,33 +291,18 @@ def main():
 
     # Display existing conversation
     for msg in st.session_state["messages"]:
-        with st.chat_message(msg["role"]):
-            st.write(msg["content"])
+        display_message(msg["role"], msg["content"])
 
     user_input = st.chat_input("Type your message here...")
     if user_input:
-        # Show user's message
-        with st.chat_message("user"):
-            st.write(user_input)
+        # Display user's message
+        display_message("user", user_input)
 
-        # Keep in session state
+        # Add user message to session state
         st.session_state["messages"].append({"role": "user", "content": user_input})
 
         # Generate assistant reply
         generate_assistant_reply(user_input)
 
-        # Track the final text from streamed tokens
-        final_text = ""
-        for msg in st.session_state.get("streamed_messages", []):
-            final_text += msg
-            
-        # Store the complete assistant response in session state
-        st.session_state["messages"].append(
-            {"role": "assistant", "content": final_text}
-        )
-        
-        # Clear streamed messages for next interaction
-        st.session_state["streamed_messages"] = []
-
 if __name__ == "__main__":
     main()