import os import re import io import time import json import queue import logging from typing import Any, Generator, Optional, List, Dict, Tuple from dataclasses import dataclass import streamlit as st from dotenv import load_dotenv from PIL import Image import openai from langsmith.wrappers import wrap_openai from langsmith import traceable # ------------------------ # Configuration and Types # ------------------------ @dataclass class AppConfig: """Application configuration settings.""" page_title: str = "Solution Specifier A" page_icon: str = "🤖" layout: str = "centered" @dataclass class Message: """Chat message structure.""" role: str content: str class StreamingError(Exception): """Custom exception for streaming-related errors.""" pass # ------------------------ # Logging Configuration # ------------------------ def setup_logging() -> logging.Logger: """Configure and return the application logger.""" logging.basicConfig( format="[%(asctime)s] %(levelname)+8s: %(message)s", level=logging.INFO, ) return logging.getLogger(__name__) logger = setup_logging() # ------------------------ # Environment Setup # ------------------------ class EnvironmentManager: """Manages environment variables and configuration.""" @staticmethod def load_environment() -> Tuple[str, str]: """Load and validate environment variables.""" load_dotenv(override=True) api_key = os.getenv("OPENAI_API_KEY") assistant_id = os.getenv("ASSISTANT_ID_SOLUTION_SPECIFIER_A") if not api_key or not assistant_id: raise RuntimeError( "Missing required environment variables. Please set " "OPENAI_API_KEY and ASSISTANT_ID_SOLUTION_SPECIFIER_A" ) return api_key, assistant_id # ------------------------ # State Management # ------------------------ class StateManager: """Manages Streamlit session state.""" @staticmethod def initialize_state() -> None: """Initialize session state variables.""" if "messages" not in st.session_state: st.session_state.messages = [] 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() if "run_stream" not in st.session_state: st.session_state.run_stream = None @staticmethod def add_message(role: str, content: str) -> None: """Add a message to the conversation history.""" st.session_state.messages.append(Message(role=role, content=content)) # ------------------------ # Text Processing # ------------------------ class TextProcessor: """Handles text processing and formatting.""" @staticmethod def remove_citations(text: str) -> str: """Remove citation markers from text.""" pattern = r"【\d+†\w+】" return re.sub(pattern, "📚", text) # ------------------------ # Streaming Handler # ------------------------ class StreamHandler: """Handles streaming of assistant responses.""" def __init__(self, client: Any): self.client = client self.text_processor = TextProcessor() self.complete_response = [] def stream_data(self) -> Generator[Any, None, None]: """Stream data from the assistant run.""" st.toast("Thinking...", icon="🤔") content_produced = False self.complete_response = [] # Reset for new stream try: for event in st.session_state.run_stream: match event.event: case "thread.message.delta": yield from self._handle_message_delta(event, content_produced) case "thread.run.requires_action": yield from self._handle_action_request(event, content_produced) case "thread.run.failed": logger.error(f"Run failed: {event}") raise StreamingError(f"Assistant run failed: {event}") st.toast("Completed", icon="✅") # Return the complete response for storage return "".join(self.complete_response) except Exception as e: logger.error(f"Streaming error: {e}") st.error(f"An error occurred while streaming: {str(e)}") raise def _handle_message_delta(self, event: Any, content_produced: bool) -> Generator[Any, None, None]: """Handle message delta events.""" content = event.data.delta.content[0] match content.type: case "text": processed_text = self.text_processor.remove_citations(content.text.value) self.complete_response.append(processed_text) # Store the chunk yield processed_text case "image_file": image_content = io.BytesIO(self.client.files.content(content.image_file.file_id).read()) yield Image.open(image_content) def _handle_action_request(self, event: Any, content_produced: bool) -> Generator[str, None, None]: """Handle action request events.""" logger.info(f"[Tool Request] {event}") st.session_state.tool_requests.put(event) if not content_produced: yield "[Processing function call...]" # ------------------------ # Tool Request Handler # ------------------------ class ToolRequestHandler: """Handles tool requests from the assistant.""" @staticmethod def handle_request(event: Any) -> Tuple[List[Dict[str, str]], str, str]: """Process tool requests and return outputs.""" st.toast("Processing function call...", icon="⚙️") tool_outputs = [] data = event.data for tool_call in data.required_action.submit_tool_outputs.tool_calls: output = ToolRequestHandler._process_tool_call(tool_call) tool_outputs.append(output) return tool_outputs, data.thread_id, data.id @staticmethod def _process_tool_call(tool_call: Any) -> Dict[str, str]: """Process individual tool calls.""" function_args = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {} match tool_call.function.name: case "hello_world": name = function_args.get("name", "anonymous") output_val = f"Hello, {name}! This was from a local function." case _: output_val = json.dumps({"status": "error", "message": "Unknown function request."}) return {"tool_call_id": tool_call.id, "output": output_val} # ------------------------ # Assistant Manager # ------------------------ class AssistantManager: """Manages interactions with the OpenAI Assistant.""" def __init__(self, client: Any, assistant_id: str): self.client = client self.assistant_id = assistant_id self.stream_handler = StreamHandler(client) self.tool_handler = ToolRequestHandler() @traceable def generate_reply(self, user_input: str) -> str: """Generate and stream assistant's reply.""" # Ensure thread exists if not st.session_state.thread: st.session_state.thread = self.client.beta.threads.create() # Add user message self.client.beta.threads.messages.create( thread_id=st.session_state.thread.id, role="user", content=user_input ) complete_response = "" # Stream initial response with self.client.beta.threads.runs.stream( thread_id=st.session_state.thread.id, assistant_id=self.assistant_id, ) as run_stream: complete_response = self._display_stream(run_stream) # Handle any tool requests self._process_tool_requests() return complete_response def _display_stream(self, run_stream: Any, create_context: bool = True) -> str: """Display streaming content.""" st.session_state.run_stream = run_stream if create_context: with st.chat_message("assistant"): return st.write_stream(self.stream_handler.stream_data) else: return st.write_stream(self.stream_handler.stream_data) def _process_tool_requests(self) -> None: """Process any pending tool requests.""" while not st.session_state.tool_requests.empty(): event = st.session_state.tool_requests.get() tool_outputs, thread_id, run_id = self.tool_handler.handle_request(event) with self.client.beta.threads.runs.submit_tool_outputs_stream( thread_id=thread_id, run_id=run_id, tool_outputs=tool_outputs ) as next_stream: self._display_stream(next_stream, create_context=False) # ------------------------ # Main Application # ------------------------ class ChatApplication: """Main chat application class.""" def __init__(self): self.config = AppConfig() api_key, assistant_id = EnvironmentManager.load_environment() # Initialize OpenAI client openai_client = openai.Client(api_key=api_key) self.client = wrap_openai(openai_client) # Initialize components self.state_manager = StateManager() self.assistant_manager = AssistantManager(self.client, assistant_id) def setup_page(self) -> None: """Configure the Streamlit page.""" st.set_page_config( page_title=self.config.page_title, page_icon=self.config.page_icon, layout=self.config.layout ) st.title(self.config.page_title) def display_chat_history(self) -> None: """Display the chat history.""" for msg in st.session_state.messages: with st.chat_message(msg.role): st.write(msg.content) def run(self) -> None: """Run the chat application.""" self.setup_page() self.state_manager.initialize_state() self.display_chat_history() user_input = st.chat_input("Type your message here...") if user_input: # Display and store user message with st.chat_message("user"): st.write(user_input) self.state_manager.add_message("user", user_input) # Generate and display assistant reply try: complete_response = self.assistant_manager.generate_reply(user_input) self.state_manager.add_message( "assistant", complete_response ) except Exception as e: st.error(f"Error generating response: {str(e)}") logger.exception("Error in assistant reply generation") def main(): """Application entry point.""" try: app = ChatApplication() app.run() except Exception as e: st.error(f"Application error: {str(e)}") logger.exception("Fatal application error") if __name__ == "__main__": main()