refactor

app.py

@@ -2,214 +2,243 @@ import os
 import re
 import streamlit as st
 from dotenv import load_dotenv
-from langsmith import traceable
-from langsmith.wrappers import wrap_openai
-import openai
-import asyncio
-import threading
-from PIL import Image
+
 import io
+import time
 import json
 import queue
 import logging
-import time
-
-# Load environment variables
-load_dotenv()
-openai.api_key = os.getenv("openai.api_key")
-LANGSMITH_API_KEY = os.getenv("LANGSMITH_API_KEY")
-ASSISTANT_ID = os.getenv("ASSISTANT_ID_SOLUTION_SPECIFIER_A")
+from PIL import Image
 
-# if not all([openai.api_key, LANGSMITH_API_KEY, ASSISTANT_ID]):
-#     raise ValueError("Please set openai.api_key, LANGSMITH_API_KEY, and ASSISTANT_ID in your .env file.")
+# ------------------------
+# LangSmith imports
+# ------------------------
+import openai
+from langsmith.wrappers import wrap_openai
+from langsmith import traceable
 
-# Initialize logging
-logging.basicConfig(format="[%(asctime)s] %(levelname)s: %(message)s", level=logging.INFO)
-logger = logging.getLogger(__name__)
+# ------------------------
+# Configure logging (optional but recommended)
+# ------------------------
+def init_logging():
+    logging.basicConfig(
+        format="[%(asctime)s] %(levelname)+8s: %(message)s",
+        level=logging.INFO,
+    )
+    return logging.getLogger()
 
-# Initialize Langsmith's traceable OpenAI client
-wrapped_openai = wrap_openai(openai.Client(api_key=openai.api_key, api_base="https://api.openai.com"))
+logger = init_logging()
 
-# Initialize Langsmith client (ensure you have configured Langsmith correctly)
-# Assuming Langsmith uses environment variables or configuration files for setup
-# If not, initialize it here accordingly
+# ------------------------
+# Load environment variables
+# ------------------------
+load_dotenv()
+api_key = os.getenv("OPENAI_API_KEY")
+assistant_id = os.getenv("ASSISTANT_ID_SOLUTION_SPECIFIER_A")  # The assistant we want to call
 
-# Define a traceable function to handle Assistant interactions
-@traceable
-def create_run(thread_id: str, assistant_id: str) -> openai.beta.RunsStream:
-    """
-    Creates a streaming run with the Assistant.
-    """
-    return wrapped_openai.beta.threads.runs.stream(
-        thread_id=thread_id,
-        assistant_id=assistant_id,
-        model="gpt-4o",  # Replace with your desired model
-        stream=True
-    )
+if not api_key or not assistant_id:
+    raise RuntimeError("Please set OPENAI_API_KEY and ASSISTANT_ID_SOLUTION_SPECIFIER_A in your environment")
 
-# Function to remove citations as per your original code
-def remove_citation(text: str) -> str:
-    pattern = r"【\d+†\w+】"
-    return re.sub(pattern, "📚", text)
+# ------------------------
+# Wrap the OpenAI client for LangSmith traceability
+# ------------------------
+openai_client = openai.Client(api_key=api_key)
+client = wrap_openai(openai_client)
 
-# Initialize session state for messages, thread_id, and tool_requests
+# ------------------------
+# Streamlit session state
+# ------------------------
 if "messages" not in st.session_state:
     st.session_state["messages"] = []
-if "thread_id" not in st.session_state:
-    st.session_state["thread_id"] = None
+
+if "thread" not in st.session_state:
+    st.session_state["thread"] = None
+
 if "tool_requests" not in st.session_state:
     st.session_state["tool_requests"] = queue.Queue()
 
 tool_requests = st.session_state["tool_requests"]
 
-# Initialize Streamlit page
-st.set_page_config(page_title="Solution Specifier A", layout="centered")
-st.title("Solution Specifier A")
+# ------------------------
+# Utility to remove citations like: 【12†somefile】
+# You can adapt to your own "annotations" handling if needed
+# ------------------------
+def remove_citation(text: str) -> str:
+    pattern = r"【\d+†\w+】"
+    return re.sub(pattern, "📚", text)
 
-# Display existing messages
-for msg in st.session_state["messages"]:
-    if msg["role"] == "user":
-        with st.chat_message("user"):
-            st.write(msg["content"])
-    else:
+# ------------------------
+# 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:")
+    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:")
+
+# ------------------------
+# 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
+    if create_context:
         with st.chat_message("assistant"):
-            if isinstance(msg["content"], Image.Image):
-                st.image(msg["content"])
-            else:
-                st.write(msg["content"])
+            streamed_result = st.write_stream(data_streamer)
+    else:
+        streamed_result = st.write_stream(data_streamer)
 
-# Chat input widget
-user_input = st.chat_input("Type your message here...")
+    # Return whatever the final token stream is
+    return streamed_result
 
-# Function to handle tool requests (function calls)
-def handle_requires_action(tool_request):
-    st.toast("Running a function", icon=":material/function:")
+# ------------------------
+# 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 = tool_request.data
-    for tool in data.required_action.submit_tool_outputs.tool_calls:
-        if tool.function.arguments:
-            function_arguments = json.loads(tool.function.arguments)
+    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_arguments = {}
-        match tool.function.name:
+            function_args = {}
+
+        match tool_call.function.name:
             case "hello_world":
-                logger.info("Calling hello_world function")
-                answer = hello_world(**function_arguments)
-                tool_outputs.append({"tool_call_id": tool.id, "output": answer})
+                # 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 _:
-                logger.error(f"Unrecognized function name: {tool.function.name}. Tool: {tool}")
-                ret_val = {
-                    "status": "error",
-                    "message": f"Function name is not recognized. Ensure the correct function name and try again."
-                }
-                tool_outputs.append({"tool_call_id": tool.id, "output": json.dumps(ret_val)})
-    st.toast("Function completed", icon=":material/function:")
+                # 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
 
-# Example function that could be called by the Assistant
-def hello_world(name: str) -> str:
-    time.sleep(2)  # Simulate a long-running task
-    return f"Hello {name}!"
-
-# Function to add assistant messages to session state
-def add_message_to_state_session(message):
-    if len(message) > 0:
-        st.session_state["messages"].append({"role": "assistant", "content": message})
-
-# Function to process streamed data
-def data_streamer(stream):
+# ------------------------
+# Main chat logic
+# ------------------------
+@traceable  # Make this function traceable via LangSmith
+def generate_assistant_reply(user_input: str):
     """
-    Stream data from the assistant. Text messages are yielded. Images and tool requests are put in the queue.
+    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.
     """
-    logger.info("Starting data streamer")
-    st.toast("Thinking...", icon=":material/emoji_objects:")
-    content_produced = False
-    try:
-        for response in stream:
-            event = response.event
-            if event == "thread.message.delta":
-                content = response.data.delta.content[0]
-                if content.type == "text":
-                    value = content.text.value
-                    content_produced = True
-                    yield value
-                elif content.type == "image_file":
-                    logger.info("Image file received")
-                    image_content = io.BytesIO(wrapped_openai.files.content(content.image_file.file_id).read())
-                    img = Image.open(image_content)
-                    content_produced = True
-                    yield img
-            elif event == "thread.run.requires_action":
-                logger.info("Run requires action")
-                tool_requests.put(response)
-                if not content_produced:
-                    yield "[LLM requires a function call]"
-                break
-            elif event == "thread.run.failed":
-                logger.error("Run failed")
-                yield "[Run failed]"
-                break
-    finally:
-        st.toast("Completed", icon=":material/emoji_objects:")
-        logger.info("Finished data streamer")
-
-# Function to display the streamed response
-def display_stream(stream):
-    with st.chat_message("assistant"):
-        for content in data_streamer(stream):
-            if isinstance(content, Image.Image):
-                st.image(content)
-                add_message_to_state_session(content)
-            else:
-                st.write(content)
-                add_message_to_state_session(content)
-
-# Main function to handle user input and assistant response
+    # Create or retrieve thread
+    if not st.session_state["thread"]:
+        st.session_state["thread"] = client.beta.threads.create()
+    thread = st.session_state["thread"]
+
+    # 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
+    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)
+
+# ------------------------
+# Streamlit UI
+# ------------------------
 def main():
+    st.set_page_config(page_title="Solution Specifier A", layout="centered")
+    st.title("Solution Specifier A")
+
+    # Display existing conversation
+    for msg in st.session_state["messages"]:
+        with st.chat_message(msg["role"]):
+            st.write(msg["content"])
+
+    user_input = st.chat_input("Type your message here...")
     if user_input:
-        # Add user message to session state
-        st.session_state["messages"].append({"role": "user", "content": user_input})
-        
-        # Display the user's message
+        # Show user's message
         with st.chat_message("user"):
             st.write(user_input)
-        
-        # Create a new thread if it doesn't exist
-        if st.session_state["thread_id"] is None:
-            logger.info("Creating new thread")
-            thread = wrapped_openai.beta.threads.create()
-            st.session_state["thread_id"] = thread.id
-        else:
-            thread = wrapped_openai.beta.threads.retrieve(st.session_state["thread_id"])
-        
-        # Add user message to the thread
-        wrapped_openai.beta.threads.messages.create(
-            thread_id=thread.id,
-            role="user",
-            content=user_input
+
+        # Keep in session state
+        st.session_state["messages"].append({"role": "user", "content": user_input})
+
+        # Generate assistant reply
+        generate_assistant_reply(user_input)
+
+        # In a real app, you might keep track of the final text
+        # from the streamed tokens. For simplicity, we store
+        # the entire streamed result as one block in session state:
+        st.session_state["messages"].append(
+            {"role": "assistant", "content": "[assistant reply streamed above]"}
         )
-        
-        # Create a new run with streaming
-        logger.info("Creating a new run with streaming")
-        stream = create_run(thread.id, ASSISTANT_ID)
-        
-        # Start a separate thread to handle streaming to avoid blocking Streamlit
-        stream_thread = threading.Thread(target=display_stream, args=(stream,))
-        stream_thread.start()
-        
-        # Handle tool requests if any
-        while not tool_requests.empty():
-            logger.info("Handling tool requests")
-            tool_request = tool_requests.get()
-            tool_outputs, thread_id, run_id = handle_requires_action(tool_request)
-            wrapped_openai.beta.threads.runs.submit_tool_outputs_stream(
-                thread_id=thread_id,
-                run_id=run_id,
-                tool_outputs=tool_outputs
-            )
-            # After handling, create a new stream to continue the conversation
-            new_stream = create_run(thread_id, ASSISTANT_ID)
-            display_stream(new_stream)
-
-# Run the main function
-main()
+
+if __name__ == "__main__":
+    main()