feat: Implement async utilities for MCP server management and JSON-RPC communication
- Added `process.py` for managing MCP server subprocesses with async capabilities. - Introduced `protocol.py` for handling JSON-RPC communication over streams. - Created `llm_client.py` to support chat completion requests to various LLM providers, integrating with MCP tools. - Defined model configurations in `llm_models.py` for different LLM providers. - Removed the synchronous `mcp_manager.py` in favor of a more modular approach. - Established a provider framework in `providers` directory with a base class and specific implementations. - Implemented `OpenAIProvider` for interacting with OpenAI's API, including streaming support and tool call handling.
This commit is contained in:
186
src/app.py
186
src/app.py
@@ -1,29 +1,111 @@
|
||||
import atexit
|
||||
import configparser
|
||||
import json # For handling potential error JSON in stream
|
||||
import logging
|
||||
|
||||
import streamlit as st
|
||||
|
||||
from openai_client import OpenAIClient
|
||||
# Updated imports
|
||||
from llm_client import LLMClient
|
||||
from src.custom_mcp.manager import SyncMCPManager # Updated import path
|
||||
|
||||
# Configure logging for the app
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def init_session_state():
|
||||
"""Initializes session state variables including clients."""
|
||||
if "messages" not in st.session_state:
|
||||
st.session_state.messages = []
|
||||
logger.info("Initialized session state: messages")
|
||||
|
||||
if "client" not in st.session_state:
|
||||
st.session_state.client = OpenAIClient()
|
||||
# Register cleanup for MCP servers
|
||||
if hasattr(st.session_state.client, "mcp_manager"):
|
||||
atexit.register(st.session_state.client.mcp_manager.shutdown)
|
||||
logger.info("Attempting to initialize clients...")
|
||||
try:
|
||||
config = configparser.ConfigParser()
|
||||
# TODO: Improve config file path handling (e.g., environment variable, absolute path)
|
||||
config_files_read = config.read("config/config.ini")
|
||||
if not config_files_read:
|
||||
raise FileNotFoundError("config.ini not found or could not be read.")
|
||||
logger.info(f"Read configuration from: {config_files_read}")
|
||||
|
||||
# --- MCP Manager Setup ---
|
||||
mcp_config_path = "config/mcp_config.json" # Default
|
||||
if config.has_section("mcp") and config["mcp"].get("servers_json"):
|
||||
mcp_config_path = config["mcp"]["servers_json"]
|
||||
logger.info(f"Using MCP config path from config.ini: {mcp_config_path}")
|
||||
else:
|
||||
logger.info(f"Using default MCP config path: {mcp_config_path}")
|
||||
|
||||
mcp_manager = SyncMCPManager(mcp_config_path)
|
||||
if not mcp_manager.initialize():
|
||||
# Log warning but continue - LLMClient will operate without tools
|
||||
logger.warning("MCP Manager failed to initialize. Proceeding without MCP tools.")
|
||||
else:
|
||||
logger.info("MCP Manager initialized successfully.")
|
||||
# Register shutdown hook for MCP manager
|
||||
atexit.register(mcp_manager.shutdown)
|
||||
logger.info("Registered MCP Manager shutdown hook.")
|
||||
|
||||
# --- LLM Client Setup ---
|
||||
# Determine provider details (e.g., from an [llm] section)
|
||||
provider_name = "openai" # Default
|
||||
model_name = None # Must be provided in config
|
||||
api_key = None
|
||||
base_url = None
|
||||
|
||||
# Prioritize [llm] section, fallback to [openai] for compatibility
|
||||
if config.has_section("llm"):
|
||||
logger.info("Reading configuration from [llm] section.")
|
||||
provider_name = config["llm"].get("provider", provider_name)
|
||||
model_name = config["llm"].get("model")
|
||||
api_key = config["llm"].get("api_key")
|
||||
base_url = config["llm"].get("base_url") # Optional
|
||||
elif config.has_section("openai"):
|
||||
logger.warning("Using legacy [openai] section for configuration.")
|
||||
provider_name = "openai" # Force openai if using this section
|
||||
model_name = config["openai"].get("model")
|
||||
api_key = config["openai"].get("api_key")
|
||||
base_url = config["openai"].get("base_url") # Optional
|
||||
else:
|
||||
raise ValueError("Missing [llm] or [openai] section in config.ini")
|
||||
|
||||
# Validate required config
|
||||
if not api_key:
|
||||
raise ValueError("Missing 'api_key' in config.ini ([llm] or [openai] section)")
|
||||
if not model_name:
|
||||
raise ValueError("Missing 'model' name in config.ini ([llm] or [openai] section)")
|
||||
|
||||
logger.info(f"Configuring LLMClient for provider: {provider_name}, model: {model_name}")
|
||||
st.session_state.client = LLMClient(
|
||||
provider_name=provider_name,
|
||||
api_key=api_key,
|
||||
mcp_manager=mcp_manager,
|
||||
base_url=base_url, # Pass None if not provided
|
||||
)
|
||||
st.session_state.model_name = model_name # Store model name for chat requests
|
||||
logger.info("LLMClient initialized successfully.")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize application clients: {e}", exc_info=True)
|
||||
st.error(f"Application Initialization Error: {e}. Please check configuration and logs.")
|
||||
# Stop the app if initialization fails critically
|
||||
st.stop()
|
||||
|
||||
|
||||
def display_chat_messages():
|
||||
"""Displays chat messages stored in session state."""
|
||||
for message in st.session_state.messages:
|
||||
with st.chat_message(message["role"]):
|
||||
# Simple markdown display for now
|
||||
st.markdown(message["content"])
|
||||
|
||||
|
||||
def handle_user_input():
|
||||
"""Handles user input, calls LLMClient, and displays the response."""
|
||||
if prompt := st.chat_input("Type your message..."):
|
||||
print(f"User input received: {prompt}") # Debug log
|
||||
logger.info(f"User input received: '{prompt[:50]}...'")
|
||||
st.session_state.messages.append({"role": "user", "content": prompt})
|
||||
with st.chat_message("user"):
|
||||
st.markdown(prompt)
|
||||
@@ -32,39 +114,85 @@ def handle_user_input():
|
||||
with st.chat_message("assistant"):
|
||||
response_placeholder = st.empty()
|
||||
full_response = ""
|
||||
error_occurred = False
|
||||
|
||||
print("Processing message...") # Debug log
|
||||
response = st.session_state.client.get_chat_response(st.session_state.messages)
|
||||
logger.info("Processing message via LLMClient...")
|
||||
# Use the new client and method, always requesting stream for UI
|
||||
response_stream = st.session_state.client.chat_completion(
|
||||
messages=st.session_state.messages,
|
||||
model=st.session_state.model_name, # Get model from session state
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Handle both MCP and standard OpenAI responses
|
||||
# Check if it's NOT a dict (assuming stream is not a dict)
|
||||
if not isinstance(response, dict):
|
||||
# Standard OpenAI streaming response
|
||||
for chunk in response:
|
||||
# Ensure chunk has choices and delta before accessing
|
||||
if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
|
||||
full_response += chunk.choices[0].delta.content
|
||||
response_placeholder.markdown(full_response + "▌")
|
||||
# Handle the response (stream generator or error dict)
|
||||
if hasattr(response_stream, "__iter__") and not isinstance(response_stream, dict):
|
||||
logger.debug("Processing response stream...")
|
||||
for chunk in response_stream:
|
||||
# Check for potential error JSON yielded by the stream
|
||||
try:
|
||||
# Attempt to parse chunk as JSON only if it looks like it
|
||||
if isinstance(chunk, str) and chunk.strip().startswith("{"):
|
||||
error_data = json.loads(chunk)
|
||||
if isinstance(error_data, dict) and "error" in error_data:
|
||||
full_response = f"Error: {error_data['error']}"
|
||||
logger.error(f"Error received in stream: {full_response}")
|
||||
st.error(full_response)
|
||||
error_occurred = True
|
||||
break # Stop processing stream on error
|
||||
# If not error JSON, treat as content chunk
|
||||
if not error_occurred and isinstance(chunk, str):
|
||||
full_response += chunk
|
||||
response_placeholder.markdown(full_response + "▌") # Add cursor effect
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# Not JSON or not error structure, treat as content chunk
|
||||
if not error_occurred and isinstance(chunk, str):
|
||||
full_response += chunk
|
||||
response_placeholder.markdown(full_response + "▌") # Add cursor effect
|
||||
|
||||
if not error_occurred:
|
||||
response_placeholder.markdown(full_response) # Final update without cursor
|
||||
logger.debug("Stream processing complete.")
|
||||
|
||||
elif isinstance(response_stream, dict) and "error" in response_stream:
|
||||
# Handle error dict returned directly (e.g., API error before streaming)
|
||||
full_response = f"Error: {response_stream['error']}"
|
||||
logger.error(f"Error returned directly from chat_completion: {full_response}")
|
||||
st.error(full_response)
|
||||
error_occurred = True
|
||||
else:
|
||||
# MCP non-streaming response
|
||||
full_response = response.get("assistant_text", "")
|
||||
response_placeholder.markdown(full_response)
|
||||
# Unexpected response type
|
||||
full_response = "[Unexpected response format from LLMClient]"
|
||||
logger.error(f"Unexpected response type: {type(response_stream)}")
|
||||
st.error(full_response)
|
||||
error_occurred = True
|
||||
|
||||
response_placeholder.markdown(full_response)
|
||||
st.session_state.messages.append({"role": "assistant", "content": full_response})
|
||||
print("Message processed successfully") # Debug log
|
||||
# Only add non-error, non-empty responses to history
|
||||
if not error_occurred and full_response:
|
||||
st.session_state.messages.append({"role": "assistant", "content": full_response})
|
||||
logger.info("Assistant response added to history.")
|
||||
elif error_occurred:
|
||||
logger.warning("Assistant response not added to history due to error.")
|
||||
else:
|
||||
logger.warning("Empty assistant response received, not added to history.")
|
||||
|
||||
except Exception as e:
|
||||
st.error(f"Error processing message: {str(e)}")
|
||||
print(f"Error details: {str(e)}") # Debug log
|
||||
logger.error(f"Error during chat handling: {str(e)}", exc_info=True)
|
||||
st.error(f"An unexpected error occurred: {str(e)}")
|
||||
|
||||
|
||||
def main():
|
||||
st.title("Streamlit Chat App")
|
||||
init_session_state()
|
||||
display_chat_messages()
|
||||
handle_user_input()
|
||||
"""Main function to run the Streamlit app."""
|
||||
st.title("MCP Chat App") # Updated title
|
||||
try:
|
||||
init_session_state()
|
||||
display_chat_messages()
|
||||
handle_user_input()
|
||||
except Exception as e:
|
||||
# Catch potential errors during rendering or handling
|
||||
logger.critical(f"Critical error in main app flow: {e}", exc_info=True)
|
||||
st.error(f"A critical application error occurred: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger.info("Starting Streamlit Chat App...")
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user