From 220e2206ff87ee469b01f4da6a7dcd0ab707503c Mon Sep 17 00:00:00 2001 From: Jagdish Parihar Date: Sun, 3 Aug 2025 20:42:29 +0530 Subject: [PATCH] Implement Model Context Protocol (MCP) for enhanced conversation context retention in the Angel Stylus Coding Assistant. Add MCPHandler for session management, update API endpoints to support session IDs, and integrate conversation history retrieval. Update .gitignore to exclude virtual environments. Add example client and testing scripts for MCP functionality. --- .gitignore | 1 + MCP_README.md | 112 +++++++++++++++++++++ QUICKSTART.md | 151 +++++++++++++++++++++++++++++ chroma_query.py | 22 ++++- llm_main_wrapper.py | 37 +++---- main.py | 43 ++++++++- mcp_api_client.py | 101 +++++++++++++++++++ mcp_handler.py | 180 ++++++++++++++++++++++++++++++++++ pages/assistant.py | 41 +++++--- rag_wrapper.py | 38 +++++++- requirements_mcp.txt | 52 ++++++++++ requirements_mcp_simple.txt | 39 ++++++++ run_mcp_assistant.py | 188 ++++++++++++++++++++++++++++++++++++ setup_database.py | 175 +++++++++++++++++++++++++++++++++ start.sh | 17 ++++ test_mcp.py | 57 +++++++++++ verify_mcp.py | 112 +++++++++++++++++++++ 17 files changed, 1328 insertions(+), 38 deletions(-) create mode 100644 MCP_README.md create mode 100644 QUICKSTART.md create mode 100644 mcp_api_client.py create mode 100644 mcp_handler.py create mode 100644 requirements_mcp.txt create mode 100644 requirements_mcp_simple.txt create mode 100755 run_mcp_assistant.py create mode 100644 setup_database.py create mode 100755 start.sh create mode 100644 test_mcp.py create mode 100644 verify_mcp.py diff --git a/.gitignore b/.gitignore index d100353..7ab925f 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ condaEnvPython/ logs/ chroma_db/ newTestPython/ +/venv \ No newline at end of file diff --git a/MCP_README.md b/MCP_README.md new file mode 100644 index 0000000..e98b386 --- /dev/null +++ b/MCP_README.md @@ -0,0 +1,112 @@ +# Model Context Protocol (MCP) Implementation for Angel Stylus Coding Assistant + +## Overview + +This document describes the Model Context Protocol (MCP) implementation in the Angel Stylus Coding Assistant. MCP enables the assistant to maintain context across multiple interactions, providing more coherent and contextually-aware responses to user queries. + +## What is MCP? + +Model Context Protocol (MCP) is a framework that allows AI models to understand and process context effectively across interactions. It ensures that models can: + +- **Remember past interactions** - Maintaining conversation history to provide consistent responses +- **Connect related information** - Understanding references to previously mentioned topics +- **Provide personalized responses** - Tailoring answers based on the user's conversation history +- **Enhance continuity** - Creating a more natural, flowing conversation experience + +## How MCP is Implemented in Angel Stylus + +### Core Components + +1. **MCPHandler Class** (`mcp_handler.py`) + - Manages conversation sessions and context storage + - Provides methods to create, retrieve, and update conversation contexts + - Handles conversation history formatting for LLM prompts + +2. **Session Management** + - Each conversation is assigned a unique session ID + - Context is maintained across multiple messages within a session + - Sessions can be created, retrieved, and updated via the API + +3. **Context Storage** + - Contexts are stored both in memory (for fast access) and on disk (for persistence) + - JSON-based storage format for easy debugging and portability + - Automatic context retrieval when continuing a conversation + +4. **API Integration** + - REST API endpoints support session-based conversations + - Session IDs can be provided by clients or generated automatically + - Response includes the session ID for future interactions + +## Usage + +### API Usage + +```python +# Example API request with session ID +import requests +import json + +# First request (no session ID) +response1 = requests.post( + "http://localhost:8001/stylus-chat", + json={ + "model": "llama3.1:8b", + "prompt": "What is Stylus?" + } +).json() + +# Get the session ID from the response +session_id = response1["session_id"] +print(f"Response: {response1['response']}") + +# Second request (with session ID) +response2 = requests.post( + "http://localhost:8001/stylus-chat", + json={ + "model": "llama3.1:8b", + "prompt": "What programming languages can I use with it?", + "session_id": session_id + } +).json() + +print(f"Response: {response2['response']}") + +# Get conversation history +history = requests.get(f"http://localhost:8001/conversation-history/{session_id}").json() +print(f"Conversation history: {history['history']}") +``` + +### Web Interface + +The Streamlit web interface automatically manages MCP sessions: + +1. Each browser session gets a unique MCP session ID +2. Conversation history is maintained as you chat +3. The "New Conversation" button in the sidebar clears the history and starts a new session + +## Testing + +To test the MCP implementation, run: + +```bash +python test_mcp.py +``` + +This script simulates a conversation with multiple turns and displays the conversation history maintained by MCP. + +## Limitations + +1. **Session Expiration**: Currently, sessions are stored indefinitely. In production, consider adding session expiration. +2. **Memory Usage**: For production use with many users, consider optimizing the in-memory cache. +3. **Token Limits**: Very long conversations may exceed the LLM's context window. The system currently limits to the 5 most recent interactions. + +## Future Improvements + +1. **Session Expiration**: Add automatic expiration for inactive sessions +2. **Advanced Context Management**: Implement smarter context selection beyond the recent message limit +3. **User Authentication**: Add user authentication to associate sessions with specific users +4. **Context Compression**: Implement techniques to compress context for more efficient storage and retrieval + +## Conclusion + +The MCP implementation enhances the Angel Stylus Coding Assistant by enabling it to maintain context across interactions, providing a more natural and helpful conversation experience. Users can now refer to previous questions and answers without needing to repeat information, making the assistant more efficient and user-friendly. \ No newline at end of file diff --git a/QUICKSTART.md b/QUICKSTART.md new file mode 100644 index 0000000..cff6f93 --- /dev/null +++ b/QUICKSTART.md @@ -0,0 +1,151 @@ +# Quick Start Guide - Angel Stylus Coding Assistant with MCP + +Get the MCP-enabled Angel Stylus Coding Assistant up and running in minutes! + +## Prerequisites + +1. **Python 3.9+** - [Download here](https://python.org/downloads/) +2. **Ollama** (optional, for LLM functionality) - [Download here](https://ollama.ai/) + +## Quick Setup (Automatic) + +### Option 1: One-Click Start (Recommended) + +**For Linux/macOS:** +```bash +./start.sh +``` + +**For Windows:** +```batch +start.bat +``` + +### Option 2: Manual Python Run +```bash +python3 run_mcp_assistant.py +``` + +This will automatically: +- ✅ Check Python version +- ✅ Install dependencies +- ✅ Set up required directories +- ✅ Initialize the ChromaDB database +- ✅ Check Ollama installation +- 🎯 Present menu options to run the assistant + +## What You'll See + +``` +🤖 Angel Stylus Coding Assistant with MCP Support +============================================================ +✅ Python version: 3.11.x +✅ Created directory: mcp_contexts +✅ Created directory: logs +✅ Created directory: chroma_db +📦 Installing dependencies... +✅ Dependencies installed successfully! +🗄️ Setting up database... +✅ Database setup completed +✅ Ollama is installed and running +Available models: + - llama3.1:8b + - deepseek-r1:7b + +============================================================ +🎯 What would you like to do? +1. Run API server (http://localhost:8001) +2. Run web interface (http://localhost:8501) +3. Run both API server and web interface +4. Test MCP functionality +5. Run API client example +6. Exit +``` + +## Usage Options + +### 1. Web Interface (Easiest) +- Choose option **2** or **3** +- Open browser to `http://localhost:8501` +- Start chatting with the MCP-enabled assistant! + +### 2. API Server +- Choose option **1** or **3** +- API available at `http://localhost:8001` +- Use the provided client examples or your own HTTP client + +### 3. Test MCP Functionality +- Choose option **4** +- Runs automated tests to verify MCP context retention works + +## First Conversation Example + +**User:** "What is Stylus?" + +**Assistant:** "Stylus is a framework for writing Arbitrum smart contracts in Rust..." + +**User:** "How do I install it?" *(Note: Assistant remembers the previous context)* + +**Assistant:** "To install Stylus (which we just discussed), you can..." + +The assistant now remembers your conversation and provides contextual responses! + +## Manual Setup (If Needed) + +If the automatic setup doesn't work: + +1. **Install dependencies:** + ```bash + pip install -r requirements_mcp.txt + ``` + +2. **Setup database:** + ```bash + python setup_database.py + ``` + +3. **Install Ollama models:** + ```bash + ollama pull llama3.1:8b + ollama pull deepseek-r1:7b + ``` + +4. **Run the assistant:** + ```bash + python main.py # For API server + # OR + streamlit run pages/assistant.py # For web interface + ``` + +## Troubleshooting + +### "Ollama not found" +- Install Ollama from https://ollama.ai/ +- Pull at least one model: `ollama pull llama3.1:8b` + +### "Module not found" errors +- Run: `pip install -r requirements_mcp.txt` +- Ensure you're using Python 3.9+ + +### Database issues +- Delete `chroma_db` folder and run setup again +- Run: `python setup_database.py` + +### Port already in use +- Change ports in the startup script +- Or kill existing processes on ports 8001/8501 + +## What Makes This Special? + +🧠 **Memory**: Unlike regular chatbots, this assistant remembers your conversation +🔄 **Context**: References to "it", "that", "the previous example" work perfectly +📚 **Knowledge**: Built on official Stylus documentation +🔧 **Flexible**: Use via web interface, API, or integrate into your own apps + +## Next Steps + +- Read [MCP_README.md](MCP_README.md) for detailed MCP implementation +- Check [README.md](README.md) for complete project documentation +- Explore the API with [mcp_api_client.py](mcp_api_client.py) + +Happy coding with Stylus! 🚀 \ No newline at end of file diff --git a/chroma_query.py b/chroma_query.py index 0ffd0e4..21ea3cb 100644 --- a/chroma_query.py +++ b/chroma_query.py @@ -3,10 +3,30 @@ from analyze_prompt import get_filters_wrapper from logger import log_info +class OllamaEmbeddingFunction: + """Custom embedding function for ChromaDB using Ollama.""" + + def __call__(self, input): + """Generate embeddings for input texts.""" + import ollama + + if isinstance(input, str): + input = [input] + + embeddings = [] + for text in input: + result = ollama.embeddings(model="nomic-embed-text", prompt=text) + embeddings.append(result["embedding"]) + return embeddings chroma_client = chromadb.PersistentClient(path="./chroma_db") -collection = chroma_client.get_or_create_collection(name="stylus_data") +# Use the same embedding function as setup_database.py +embedding_function = OllamaEmbeddingFunction() +collection = chroma_client.get_or_create_collection( + name="stylus_data", + embedding_function=embedding_function +) def get_prompt_embedding(user_prompt): return ollama.embeddings(model="nomic-embed-text", prompt=user_prompt)["embedding"] diff --git a/llm_main_wrapper.py b/llm_main_wrapper.py index bd659af..c449317 100644 --- a/llm_main_wrapper.py +++ b/llm_main_wrapper.py @@ -10,7 +10,7 @@ def join_chunks_limited(chunks, max_chars=10000): combined += chunk + "\n\n" return combined.strip() -def stylus_request_with_llm(model, user_prompt): +def stylus_request_with_llm(model, user_prompt, conversation_history=""): docs = get_chroma_documents(user_prompt) if not docs: @@ -20,11 +20,13 @@ def stylus_request_with_llm(model, user_prompt): The user asked: "{user_prompt}" No relevant content was retrieved for this question — either because: - – The docs don’t cover this topic yet + – The docs don't cover this topic yet – The question was too vague – Or it's outside the scope of Stylus (e.g. generic Solidity, unrelated tooling, etc.) - ⚠️ You are not keeping track of any previous messages. If the user's question seems to refer to something from earlier, let them know you don't have access to that context and ask them to clarify. + ⚠️ You are keeping track of previous messages using MCP (Model Context Protocol). If the user's question seems to reference something from earlier, check the conversation history provided below. + + {conversation_history} Keep your response short and direct. @@ -33,12 +35,13 @@ def stylus_request_with_llm(model, user_prompt): – Which tool they're using (Rust SDK, Rust CLI, etc.) – The actual command or error they're dealing with - If they already provided enough detail and we just don’t have docs for it yet, tell them that clearly. Suggest they rephrase or check back later — the docs are still growing. + If they already provided enough detail and we just don't have docs for it yet, tell them that clearly. Suggest they rephrase or check back later — the docs are still growing. ⚠️ Do **not** make anything up. If it's not in the docs, just say that. """ - return call_llm(fallback_prompt, user_prompt, model) + response = call_llm(fallback_prompt, user_prompt, model) + return response, [] context = join_chunks_limited(docs) @@ -46,12 +49,20 @@ def stylus_request_with_llm(model, user_prompt): You are a developer assistant for Stylus (Arbitrum), helping users by answering technical questions based strictly on the official documentation. This is a Retrieval-Augmented Generation (RAG) system. The information provided below was automatically retrieved from the official Stylus documentation, based on the user's question. - + Only use the information in the context below. Do **not** rely on any prior knowledge or external sources. If the context includes a URL, you may include it in your response — otherwise, never guess or generate links. ⚠️ Important: + - You are using MCP (Model Context Protocol) to maintain conversation context across interactions + - Previous conversation history is provided below, use it to provide coherent responses + - If the user refers to something from a previous message, check the conversation history + + --- CONVERSATION HISTORY --- + {conversation_history} + --- END OF CONVERSATION HISTORY --- + If the context doesn't contain the necessary information to answer the question, say: - "I'm sorry, I couldn’t find specific information to help with that right now. The docs are still evolving — feel free to check back later." + "I'm sorry, I couldn't find specific information to help with that right now. The docs are still evolving — feel free to check back later." Your tone should be direct, clear, and practical — like a developer helping another developer. No fluff, no guessing. @@ -60,15 +71,5 @@ def stylus_request_with_llm(model, user_prompt): --- END OF CONTEXT --- """ - - - #response = call_llm(formatted_prompt, user_prompt,"qwen2.5:32b") - #log_info("calling llm now") response = call_llm(formatted_prompt, user_prompt, model) - return response - - -#print(plan_trip_with_llm("what information is available for traveling with a dog to France namely about Documents?")) -#print(plan_trip_with_llm("I'll be travelling to France with my cat what usefull tips can you give me?")) -#print(plan_trip_with_llm("Can I travel by ferry from uk to france with my cat?")) -#print(plan_trip_with_llm("I will travel to france with my cat is there any concerns regarding border stuff?")) \ No newline at end of file + return response, docs \ No newline at end of file diff --git a/main.py b/main.py index 7b6336b..add9921 100644 --- a/main.py +++ b/main.py @@ -6,24 +6,61 @@ from typing import Optional from logger import log_info, log_separator import time +import uuid +from mcp_handler import MCPHandler app = FastAPI() +mcp_handler = MCPHandler() class StylusRequest(BaseModel): model: str prompt: str - + session_id: Optional[str] = None @app.post("/stylus-chat") def stylus_chat(request: StylusRequest): log_info(f"User started a request | Model: {request.model} | Prompt: {request.prompt}") + + # Generate a session ID if not provided + session_id = request.session_id or str(uuid.uuid4()) + + # Get conversation history for MCP context + conversation_history = mcp_handler.get_conversation_history(session_id) + start_time = time.time() - result = stylus_request_with_llm(request.model, request.prompt) + + # Pass the conversation history to the LLM wrapper + result, retrieved_docs = stylus_request_with_llm( + request.model, + request.prompt, + conversation_history + ) + clean_result = remove_think_tags(result) + + # Store the interaction in MCP + mcp_handler.add_interaction( + session_id, + request.prompt, + clean_result, + retrieved_docs + ) + duration = round(time.time() - start_time, 2) log_info(f"✅ Finished inference | Time: {duration}s | Output: {clean_result[:50]}...") log_separator() - return {"response": clean_result} + + # Return the session ID along with the response + return { + "response": clean_result, + "session_id": session_id + } + +# Add an endpoint to get conversation history +@app.get("/conversation-history/{session_id}") +def get_conversation_history(session_id: str): + history = mcp_handler.get_conversation_history(session_id, max_interactions=10) + return {"history": history, "session_id": session_id} if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=8001) diff --git a/mcp_api_client.py b/mcp_api_client.py new file mode 100644 index 0000000..ae76033 --- /dev/null +++ b/mcp_api_client.py @@ -0,0 +1,101 @@ +""" +Example client script for using the MCP-enabled Angel Stylus Coding Assistant API. +""" + +import requests +import json +import time + +# API endpoint +API_BASE_URL = "http://localhost:8001" + +def query_assistant(prompt, model="llama3.1:8b", session_id=None): + """ + Query the assistant with optional session ID for context preservation. + + Args: + prompt: User's question + model: LLM model to use + session_id: Optional session ID to continue a conversation + + Returns: + Response text and session ID + """ + request_data = { + "prompt": prompt, + "model": model + } + + # Include session ID if provided + if session_id: + request_data["session_id"] = session_id + + response = requests.post( + f"{API_BASE_URL}/stylus-chat", + json=request_data + ) + + if response.status_code != 200: + print(f"Error: {response.status_code} - {response.text}") + return None, None + + response_data = response.json() + return response_data["response"], response_data["session_id"] + +def get_conversation_history(session_id): + """ + Retrieve conversation history for a session. + + Args: + session_id: Session ID + + Returns: + Conversation history as a string + """ + response = requests.get(f"{API_BASE_URL}/conversation-history/{session_id}") + + if response.status_code != 200: + print(f"Error: {response.status_code} - {response.text}") + return None + + return response.json()["history"] + +def run_example_conversation(): + """ + Run an example conversation to demonstrate MCP capabilities. + """ + print("Angel Stylus Coding Assistant with MCP\n") + + # First query (no session ID) + print("User: What is Stylus?") + response1, session_id = query_assistant("What is Stylus?") + print(f"Assistant: {response1}\n") + print(f"Session ID: {session_id}\n") + + time.sleep(1) # Small delay for readability + + # Second query (with session ID) + print("User: How can I use it with Rust?") + response2, session_id = query_assistant("How can I use it with Rust?", session_id=session_id) + print(f"Assistant: {response2}\n") + + time.sleep(1) # Small delay for readability + + # Third query (follow-up with context) + print("User: Show me an example contract") + response3, session_id = query_assistant("Show me an example contract", session_id=session_id) + print(f"Assistant: {response3}\n") + + # Display the full conversation history + print("\n=== Conversation History ===") + history = get_conversation_history(session_id) + print(history) + +if __name__ == "__main__": + print("Starting MCP API client example...\n") + try: + run_example_conversation() + except Exception as e: + print(f"Error: {str(e)}") + print("Make sure the API server is running at http://localhost:8001") + print("\nDone.") \ No newline at end of file diff --git a/mcp_handler.py b/mcp_handler.py new file mode 100644 index 0000000..84bd88e --- /dev/null +++ b/mcp_handler.py @@ -0,0 +1,180 @@ +import json +import os +import time +from typing import Dict, List, Optional, Any +from logger import log_info + +class MCPHandler: + """ + Model Context Protocol (MCP) handler for maintaining conversation context + across multiple interactions with the Angel Stylus Coding Assistant. + """ + + def __init__(self, context_file_path="./mcp_contexts"): + """ + Initialize the MCP handler with the path to store context files. + + Args: + context_file_path: Directory to store context files + """ + self.context_file_path = context_file_path + self.active_contexts = {} # In-memory cache of active contexts + + # Ensure context directory exists + if not os.path.exists(context_file_path): + os.makedirs(context_file_path) + + def create_session(self, session_id: str, initial_data: Optional[Dict] = None) -> Dict: + """ + Create a new MCP session with optional initial data. + + Args: + session_id: Unique identifier for the session + initial_data: Optional initial context data + + Returns: + The created context + """ + context = { + "session_id": session_id, + "created_at": time.time(), + "last_updated": time.time(), + "interactions": [], + "metadata": initial_data or {} + } + + # Save to disk and memory + self._save_context(session_id, context) + self.active_contexts[session_id] = context + + log_info(f"Created new MCP session: {session_id}") + return context + + def add_interaction(self, session_id: str, user_prompt: str, system_response: str, + retrieved_docs: Optional[List[str]] = None) -> Dict: + """ + Add a new interaction to an existing session. + + Args: + session_id: Session ID + user_prompt: The user's input + system_response: The system's response + retrieved_docs: Optional list of retrieved document chunks + + Returns: + Updated context + """ + # Get or create context + context = self.get_context(session_id) + if not context: + context = self.create_session(session_id) + + # Add the new interaction + interaction = { + "timestamp": time.time(), + "user_prompt": user_prompt, + "system_response": system_response, + "retrieved_docs": retrieved_docs or [] + } + + context["interactions"].append(interaction) + context["last_updated"] = time.time() + + # Save updated context + self._save_context(session_id, context) + + return context + + def get_context(self, session_id: str) -> Optional[Dict]: + """ + Retrieve the context for a session. + + Args: + session_id: Session ID + + Returns: + Context data or None if not found + """ + # Check in-memory cache first + if session_id in self.active_contexts: + return self.active_contexts[session_id] + + # Try to load from disk + context_file = os.path.join(self.context_file_path, f"{session_id}.json") + if os.path.exists(context_file): + try: + with open(context_file, 'r') as f: + context = json.load(f) + self.active_contexts[session_id] = context + return context + except Exception as e: + log_info(f"Error loading context for session {session_id}: {str(e)}") + + return None + + def get_conversation_history(self, session_id: str, max_interactions: int = 5) -> str: + """ + Get formatted conversation history for a session, limited to the most recent interactions. + + Args: + session_id: Session ID + max_interactions: Maximum number of past interactions to include + + Returns: + Formatted conversation history + """ + context = self.get_context(session_id) + if not context or not context.get("interactions"): + return "" + + # Get the most recent interactions up to max_interactions + recent_interactions = context["interactions"][-max_interactions:] + + # Format the conversation history + history = [] + for interaction in recent_interactions: + history.append(f"User: {interaction['user_prompt']}") + history.append(f"Assistant: {interaction['system_response']}") + + return "\n\n".join(history) + + def _save_context(self, session_id: str, context: Dict) -> None: + """ + Save context to disk. + + Args: + session_id: Session ID + context: Context data to save + """ + context_file = os.path.join(self.context_file_path, f"{session_id}.json") + try: + with open(context_file, 'w') as f: + json.dump(context, f, indent=2) + + # Update in-memory cache + self.active_contexts[session_id] = context + except Exception as e: + log_info(f"Error saving context for session {session_id}: {str(e)}") + + def get_relevant_docs_from_history(self, session_id: str) -> List[str]: + """ + Get a list of relevant document chunks from past interactions. + + Args: + session_id: Session ID + + Returns: + List of relevant document chunks + """ + context = self.get_context(session_id) + if not context or not context.get("interactions"): + return [] + + # Collect unique document chunks from past interactions + unique_docs = set() + for interaction in context["interactions"]: + if interaction.get("retrieved_docs"): + for doc in interaction["retrieved_docs"]: + unique_docs.add(doc) + + return list(unique_docs) \ No newline at end of file diff --git a/pages/assistant.py b/pages/assistant.py index 1db0863..98da71b 100755 --- a/pages/assistant.py +++ b/pages/assistant.py @@ -1,25 +1,45 @@ -from rag_query import query_rag -from rag_wrapper import stylus_chat import streamlit as st +import sys +import os + +# Add project root to path for imports +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +# Now import from the parent directory +from rag_wrapper import stylus_chat +from mcp_handler import MCPHandler + +# Initialize the MCP handler +mcp_handler = MCPHandler() with st.sidebar: # Reset Vector DAta option = st.selectbox( "Model to use", - ("deepseek-r1:7b", "llama3.1:8b", "deepseek-r1:14b", "qwen2.5:32b") - #("mistral", "llama3", "llama3.3", "deepseek-r1:7b"), + ("llama3.1:8b", "mistral") ) st.write("You selected:", option) + + # Add option to clear conversation history + if st.button("New Conversation"): + st.session_state.messages = [{"role": "assistant", "content": "Hi! I'm your Stylus coding assistant with MCP context retention. Ask me anything about Stylus development!"}] + st.session_state.session_id = None + +# Initialize session ID if not present +if "session_id" not in st.session_state: + st.session_state.session_id = None # Function for generating LLM response -def generate_response(input): - result = stylus_chat(option, input) +def generate_response(input, session_id=None): + result, session_id = stylus_chat(option, input, session_id) + # Store the session ID for future use + st.session_state.session_id = session_id return result # Store LLM generated responses if "messages" not in st.session_state.keys(): - st.session_state.messages = [{"role": "assistant", "content": "Hi! I'm your Stylus coding assistant. Please note that I will not remember previous messages, so include all the details you need help with in each question. How can I help today?"}] + st.session_state.messages = [{"role": "assistant", "content": "Hi! I'm your Stylus coding assistant with MCP context retention. Ask me anything about Stylus development!"}] # Display chat messages for message in st.session_state.messages: @@ -35,10 +55,9 @@ def generate_response(input): # Generate a new response if last message is not from assistant if st.session_state.messages[-1]["role"] != "assistant": with st.chat_message("assistant"): - with st.spinner("Getting your answer from superior intelligence.."): - - #response = generate_response("".join([ "".join([dict["role"], dict["content"]]) for dict in st.session_state.messages])) - response = generate_response(input) + with st.spinner("Getting your answer using MCP context..."): + # Pass the session_id to maintain context across interactions + response = generate_response(input, st.session_state.session_id) st.write(response) message = {"role": "assistant", "content": response} st.session_state.messages.append(message) \ No newline at end of file diff --git a/rag_wrapper.py b/rag_wrapper.py index 6cd4f07..d35343f 100644 --- a/rag_wrapper.py +++ b/rag_wrapper.py @@ -2,16 +2,44 @@ from logger import log_info, log_separator import time from aux_functions import remove_think_tags +import uuid +from mcp_handler import MCPHandler +# Initialize the MCP handler +mcp_handler = MCPHandler() -def stylus_chat(model_name: str, prompt: str): - log_info(f"User started a request | Model: {model_name} | Prompt: {prompt}") +def stylus_chat(model_name: str, prompt: str, session_id=None): + # Generate a session ID if not provided + if not session_id: + session_id = str(uuid.uuid4()) + log_info(f"Created new MCP session: {session_id}") + else: + log_info(f"Using existing MCP session: {session_id}") + + # Get conversation history + conversation_history = mcp_handler.get_conversation_history(session_id) + + log_info(f"User started a request | Model: {model_name} | Prompt: {prompt} | Session: {session_id}") start_time = time.time() - result = stylus_request_with_llm(model_name, prompt) + + # Pass the conversation history to the LLM wrapper + result, retrieved_docs = stylus_request_with_llm(model_name, prompt, conversation_history) + clean_result = remove_think_tags(result) + + # Store the interaction in MCP + mcp_handler.add_interaction( + session_id, + prompt, + clean_result, + retrieved_docs + ) + duration = round(time.time() - start_time, 2) - log_info(f"✅ Finished inference | Time: {duration}s | Output: {clean_result[:50]}...") + log_info(f"✅ Finished inference | Time: {duration}s | Output: {clean_result[:50]}... | Session: {session_id}") log_separator() - return clean_result + + # Return both the response and the session ID + return clean_result, session_id diff --git a/requirements_mcp.txt b/requirements_mcp.txt new file mode 100644 index 0000000..7c9172c --- /dev/null +++ b/requirements_mcp.txt @@ -0,0 +1,52 @@ +# Core dependencies for Angel Stylus Coding Assistant with MCP support + +# Web Framework +fastapi==0.115.12 +uvicorn==0.34.0 +streamlit==1.41.1 + +# HTTP Client +requests==2.32.3 +httpx==0.28.1 + +# Vector Database and Embeddings +chromadb==0.6.3 +chroma-hnswlib==0.7.6 + +# LangChain for RAG (compatible versions) +langchain-community==0.2.16 +langchain>=0.2.16,<0.3.0 +langchain-text-splitters>=0.2.0,<0.3.0 + +# LLM Integration +ollama==0.4.9 +litellm>=1.0.0 + +# Document Processing +pypdf==5.2.0 +python-multipart==0.0.18 + +# Data Processing +pandas==2.2.3 +numpy==2.0.2 +fuzzywuzzy==0.18.0 +python-levenshtein==0.26.1 + +# JSON and Data Utils +jq==1.8.0 +pydantic==2.10.5 + +# Utilities +python-dotenv==1.0.1 +coloredlogs==15.0.1 +tqdm==4.67.1 + +# Additional dependencies that might be needed +aiohttp==3.11.14 +anyio==4.9.0 +backoff==2.2.1 +diskcache==5.6.3 +filelock==3.18.0 +gitpython==3.1.44 +humanfriendly==10.0 +tenacity==9.0.0 \ No newline at end of file diff --git a/requirements_mcp_simple.txt b/requirements_mcp_simple.txt new file mode 100644 index 0000000..5a9e4d2 --- /dev/null +++ b/requirements_mcp_simple.txt @@ -0,0 +1,39 @@ +# Essential dependencies for Angel Stylus Coding Assistant with MCP support +# Using flexible versions to avoid conflicts + +# Web Framework +fastapi +uvicorn +streamlit + +# HTTP Client +requests +httpx + +# Vector Database and Embeddings +chromadb +chroma-hnswlib + +# LangChain for RAG (letting pip resolve versions) +langchain-community +langchain +langchain-text-splitters + +# LLM Integration +ollama +litellm + +# Data Processing (flexible versions) +pandas +numpy +fuzzywuzzy +python-levenshtein + +# JSON and Data Utils +jq +pydantic + +# Utilities +python-dotenv +coloredlogs +tqdm \ No newline at end of file diff --git a/run_mcp_assistant.py b/run_mcp_assistant.py new file mode 100755 index 0000000..a339480 --- /dev/null +++ b/run_mcp_assistant.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +""" +Startup script for the MCP-enabled Angel Stylus Coding Assistant. +This script handles setup, dependency checks, and provides options to run different components. +""" + +import subprocess +import sys +import os +import time +import threading +from pathlib import Path + +def check_python_version(): + """Check if Python version is compatible.""" + if sys.version_info < (3, 9): + print("❌ Python 3.9 or higher is required!") + print(f"Current version: {sys.version}") + return False + print(f"✅ Python version: {sys.version}") + return True + +def install_dependencies(): + """Install required dependencies.""" + print("📦 Installing dependencies...") + try: + subprocess.check_call([ + sys.executable, "-m", "pip", "install", "-r", "requirements_mcp.txt" + ]) + print("✅ Dependencies installed successfully!") + return True + except subprocess.CalledProcessError as e: + print(f"❌ Error installing dependencies: {e}") + return False + +def check_ollama(): + """Check if Ollama is installed and running.""" + try: + result = subprocess.run(["ollama", "list"], capture_output=True, text=True) + if result.returncode == 0: + print("✅ Ollama is installed and running") + print("Available models:") + for line in result.stdout.strip().split('\n')[1:]: # Skip header + if line.strip(): + model_name = line.split()[0] + print(f" - {model_name}") + return True + else: + print("⚠️ Ollama is installed but not running") + return False + except FileNotFoundError: + print("❌ Ollama is not installed!") + print("Please install Ollama from: https://ollama.ai/") + return False + +def setup_directories(): + """Create necessary directories.""" + directories = ["mcp_contexts", "logs", "chroma_db"] + for directory in directories: + Path(directory).mkdir(exist_ok=True) + print(f"✅ Created directory: {directory}") + +def setup_database(): + """Setup the ChromaDB database.""" + print("🗄️ Setting up database...") + try: + subprocess.run([sys.executable, "setup_database.py"]) + print("✅ Database setup completed") + except Exception as e: + print(f"⚠️ Database setup issue: {e}") + print("You can run setup manually with: python setup_database.py") + +def run_api_server(): + """Run the FastAPI server.""" + print("🚀 Starting API server...") + try: + subprocess.run([sys.executable, "main.py"]) + except KeyboardInterrupt: + print("\n🛑 API server stopped") + +def run_web_interface(): + """Run the Streamlit web interface.""" + print("🌐 Starting web interface...") + try: + subprocess.run([ + sys.executable, "-m", "streamlit", "run", + "pages/assistant.py", "--server.port", "8501" + ]) + except KeyboardInterrupt: + print("\n🛑 Web interface stopped") + +def run_test(): + """Run the MCP test.""" + print("🧪 Running MCP test...") + try: + subprocess.run([sys.executable, "test_mcp.py"]) + except Exception as e: + print(f"❌ Test failed: {e}") + +def run_api_client_example(): + """Run the API client example.""" + print("📡 Running API client example...") + try: + subprocess.run([sys.executable, "mcp_api_client.py"]) + except Exception as e: + print(f"❌ API client example failed: {e}") + +def main(): + """Main function to coordinate the startup process.""" + print("🤖 Angel Stylus Coding Assistant with MCP Support") + print("=" * 60) + + # Check Python version + if not check_python_version(): + return + + # Setup directories + setup_directories() + + # Check if requirements file exists + if not os.path.exists("requirements_mcp.txt"): + print("❌ requirements_mcp.txt not found!") + return + + # Install dependencies + install_dependencies() + + # Setup database if needed + setup_database() + + # Check Ollama + ollama_ok = check_ollama() + if not ollama_ok: + print("\n⚠️ Ollama is required for the LLM functionality.") + print("You can still test the MCP framework, but LLM responses won't work.") + + print("\n" + "=" * 60) + print("🎯 What would you like to do?") + print("1. Run API server (http://localhost:8001)") + print("2. Run web interface (http://localhost:8501)") + print("3. Run both API server and web interface") + print("4. Test MCP functionality") + print("5. Run API client example") + print("6. Exit") + + while True: + try: + choice = input("\nEnter your choice (1-6): ").strip() + + if choice == "1": + run_api_server() + break + elif choice == "2": + run_web_interface() + break + elif choice == "3": + print("🚀 Starting both API server and web interface...") + print("API server will run on http://localhost:8001") + print("Web interface will run on http://localhost:8501") + print("Press Ctrl+C to stop both services") + + # Start API server in a separate thread + api_thread = threading.Thread(target=run_api_server, daemon=True) + api_thread.start() + + # Wait a moment for API server to start + time.sleep(3) + + # Start web interface in main thread + run_web_interface() + break + elif choice == "4": + run_test() + break + elif choice == "5": + run_api_client_example() + break + elif choice == "6": + print("👋 Goodbye!") + break + else: + print("❌ Invalid choice. Please enter 1-6.") + except KeyboardInterrupt: + print("\n👋 Goodbye!") + break + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/setup_database.py b/setup_database.py new file mode 100644 index 0000000..308e776 --- /dev/null +++ b/setup_database.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +""" +Setup script to initialize the ChromaDB with Stylus documentation data. +""" + +import json +import os +import sys +from pathlib import Path + +class OllamaEmbeddingFunction: + """Custom embedding function for ChromaDB using Ollama.""" + + def __call__(self, input): + """Generate embeddings for input texts.""" + import ollama + + if isinstance(input, str): + input = [input] + + embeddings = [] + for text in input: + result = ollama.embeddings(model="nomic-embed-text", prompt=text) + embeddings.append(result["embedding"]) + return embeddings + +def setup_database(): + """Initialize the ChromaDB with available data.""" + print("🗄️ Setting up ChromaDB with Stylus documentation...") + + # Create chroma_db directory if it doesn't exist + Path("chroma_db").mkdir(exist_ok=True) + + # Check if we have the necessary import modules + try: + import chromadb + import ollama + print("✅ Required modules found") + except ImportError as e: + print(f"❌ Missing required module: {e}") + print("Please install dependencies first: pip install -r requirements_mcp.txt") + return False + + try: + # Initialize ChromaDB client + client = chromadb.PersistentClient(path="./chroma_db") + + # Create or get collection with custom embedding function + embedding_function = OllamaEmbeddingFunction() + collection = client.get_or_create_collection( + name="stylus_data", + embedding_function=embedding_function + ) + + # Check if collection already has data + count = collection.count() + if count > 0: + print(f"✅ Database already initialized with {count} documents") + return True + + # Load data from JSON files + data_files = [ + "data/stylus_docs.json", + "data/arbitrum-stylus-data.json", + "data/stylus-dataset.json" + ] + + documents = [] + metadatas = [] + ids = [] + + doc_id = 0 + + for file_path in data_files: + if os.path.exists(file_path): + print(f"📖 Loading data from {file_path}...") + + try: + with open(file_path, 'r', encoding='utf-8') as f: + data = json.load(f) + + # Handle different JSON structures + if isinstance(data, list): + items = data + elif isinstance(data, dict) and 'documents' in data: + items = data['documents'] + else: + items = [data] + + for item in items: + if isinstance(item, dict): + # Extract text content + text_content = "" + if 'content' in item: + text_content = str(item['content']) + elif 'text' in item: + text_content = str(item['text']) + elif 'page_content' in item: + text_content = str(item['page_content']) + else: + text_content = str(item) + + if text_content and len(text_content.strip()) > 10: + documents.append(text_content) + + # Create metadata (ensure all values are strings) + metadata = { + "source": file_path, + "doc_type": "stylus_documentation" + } + + # Add any additional metadata from the item (convert to strings) + if 'metadata' in item and isinstance(item['metadata'], dict): + for key, value in item['metadata'].items(): + if isinstance(value, (str, int, float, bool)): + metadata[key] = str(value) + elif isinstance(value, list): + metadata[key] = ", ".join(str(v) for v in value) + else: + metadata[key] = str(value) + + metadatas.append(metadata) + ids.append(f"doc_{doc_id}") + doc_id += 1 + + except Exception as e: + print(f"⚠️ Error loading {file_path}: {e}") + continue + else: + print(f"⚠️ File not found: {file_path}") + + if documents: + print(f"📚 Adding {len(documents)} documents to ChromaDB...") + + # Add documents in batches to avoid memory issues + batch_size = 100 + for i in range(0, len(documents), batch_size): + batch_end = min(i + batch_size, len(documents)) + batch_docs = documents[i:batch_end] + batch_metadata = metadatas[i:batch_end] + batch_ids = ids[i:batch_end] + + collection.add( + documents=batch_docs, + metadatas=batch_metadata, + ids=batch_ids + ) + print(f" Added batch {i//batch_size + 1}/{(len(documents)-1)//batch_size + 1}") + + print(f"✅ Successfully added {len(documents)} documents to ChromaDB") + return True + else: + print("❌ No valid documents found in data files") + return False + + except Exception as e: + print(f"❌ Error setting up database: {e}") + return False + +def main(): + """Main function.""" + print("🤖 Angel Stylus Coding Assistant - Database Setup") + print("=" * 50) + + success = setup_database() + + if success: + print("\n✅ Database setup completed successfully!") + print("You can now run the assistant with: python run_mcp_assistant.py") + else: + print("\n❌ Database setup failed!") + print("Please check the error messages above and try again.") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/start.sh b/start.sh new file mode 100755 index 0000000..9d53a24 --- /dev/null +++ b/start.sh @@ -0,0 +1,17 @@ +#!/bin/bash + +# Angel Stylus Coding Assistant with MCP - Startup Script +echo "🤖 Angel Stylus Coding Assistant with MCP Support" +echo "==================================================" + +# Check if Python is available +if ! command -v python3 &> /dev/null; then + echo "❌ Python 3 is not installed!" + exit 1 +fi + +# Make the Python script executable +chmod +x run_mcp_assistant.py + +# Run the Python startup script +python3 run_mcp_assistant.py \ No newline at end of file diff --git a/test_mcp.py b/test_mcp.py new file mode 100644 index 0000000..14e40ab --- /dev/null +++ b/test_mcp.py @@ -0,0 +1,57 @@ +""" +Test script for the MCP (Model Context Protocol) implementation. +This script simulates a conversation with the Stylus assistant using MCP. +""" + +import time +import uuid +from mcp_handler import MCPHandler +from rag_wrapper import stylus_chat +from logger import log_info, log_separator + +def test_mcp_conversation(): + """ + Test a multi-turn conversation with MCP support. + """ + # Initialize the MCP handler + mcp = MCPHandler() + + # Generate a session ID for this conversation + session_id = str(uuid.uuid4()) + print(f"Starting conversation with session ID: {session_id}") + + # First message + prompt1 = "What is Stylus?" + print(f"\nUser: {prompt1}") + response1, session_id = stylus_chat("llama3.1:8b", prompt1, session_id) + print(f"Assistant: {response1}") + + time.sleep(1) # Small delay for readability + + # Second message with context + prompt2 = "What programming languages can I use with it?" + print(f"\nUser: {prompt2}") + response2, session_id = stylus_chat("llama3.1:8b", prompt2, session_id) + print(f"Assistant: {response2}") + + time.sleep(1) # Small delay for readability + + # Third message asking for code + prompt3 = "Give me an example of a simple Rust program for Stylus." + print(f"\nUser: {prompt3}") + response3, session_id = stylus_chat("llama3.1:8b", prompt3, session_id) + print(f"Assistant: {response3}") + + # Display the conversation history from MCP + history = mcp.get_conversation_history(session_id) + print("\n--- Conversation History from MCP ---") + print(history) + + return session_id + +if __name__ == "__main__": + log_separator() + log_info("Starting MCP test") + session_id = test_mcp_conversation() + log_info(f"MCP test completed with session ID: {session_id}") + log_separator() \ No newline at end of file diff --git a/verify_mcp.py b/verify_mcp.py new file mode 100644 index 0000000..9cc9d8d --- /dev/null +++ b/verify_mcp.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +""" +Verification script to test MCP functionality +""" + +import requests +import json +import time + +def test_mcp_api(): + """Test the MCP-enabled API""" + print("🧪 Testing MCP-enabled Angel Stylus Coding Assistant") + print("=" * 60) + + api_url = "http://localhost:8001/stylus-chat" + + # Test 1: First request + print("\n📤 Test 1: Initial request") + payload1 = { + "model": "mistral", + "prompt": "Hello, can you help me with programming?" + } + + try: + response1 = requests.post(api_url, json=payload1, timeout=30) + if response1.status_code == 200: + data1 = response1.json() + session_id = data1.get("session_id") + print(f"✅ Response received") + print(f"📝 Assistant: {data1['response'][:100]}...") + print(f"🔑 Session ID: {session_id}") + + # Test 2: Follow-up request with context + print("\n📤 Test 2: Follow-up request (testing MCP context)") + payload2 = { + "model": "mistral", + "prompt": "What did I just ask you about?", + "session_id": session_id + } + + time.sleep(1) # Small delay + response2 = requests.post(api_url, json=payload2, timeout=30) + + if response2.status_code == 200: + data2 = response2.json() + print(f"✅ Follow-up response received") + print(f"📝 Assistant: {data2['response'][:100]}...") + print(f"🔑 Same session ID: {data2.get('session_id') == session_id}") + + # Check if context was retained + if "programming" in data2['response'].lower() or "help" in data2['response'].lower(): + print("🎉 MCP CONTEXT RETENTION: SUCCESS!") + print(" The assistant remembered our previous conversation!") + else: + print("⚠️ MCP context retention may not be working as expected") + + # Test 3: Get conversation history + print("\n📤 Test 3: Get conversation history") + history_url = f"http://localhost:8001/conversation-history/{session_id}" + history_response = requests.get(history_url, timeout=10) + + if history_response.status_code == 200: + history_data = history_response.json() + print("✅ Conversation history retrieved") + print(f"📚 History preview: {history_data['history'][:150]}...") + else: + print(f"❌ Failed to get history: {history_response.status_code}") + + else: + print(f"❌ Follow-up request failed: {response2.status_code}") + print(response2.text) + else: + print(f"❌ Initial request failed: {response1.status_code}") + print(response1.text) + + except requests.exceptions.ConnectionError: + print("❌ Connection failed. Make sure the API server is running on http://localhost:8001") + print(" Run: python main.py") + except Exception as e: + print(f"❌ Test failed: {str(e)}") + +def check_services(): + """Check if required services are running""" + print("\n🔍 Checking services...") + + # Check API server + try: + response = requests.get("http://localhost:8001", timeout=5) + print("✅ API server is running on http://localhost:8001") + except: + print("❌ API server not accessible on http://localhost:8001") + print(" Start with: python main.py") + + # Check Streamlit (if running) + try: + response = requests.get("http://localhost:8501", timeout=5) + print("✅ Streamlit web interface is running on http://localhost:8501") + except: + print("ℹ️ Streamlit web interface not running on http://localhost:8501") + print(" Start with: streamlit run pages/assistant.py") + +if __name__ == "__main__": + check_services() + test_mcp_api() + + print("\n" + "=" * 60) + print("🎉 MCP verification complete!") + print("✨ Your Angel Stylus Coding Assistant with MCP is ready!") + print("\n🌐 Access methods:") + print(" • API: http://localhost:8001") + print(" • Web UI: http://localhost:8501") + print(" • Test again: python verify_mcp.py") \ No newline at end of file