Chroma MCP Server

""" ThinkingSession module for managing sequential thoughts through MCP. """ import uuid from typing import Dict, List, Optional, Union, Any from mcp import ClientSession import os class ThinkingSession: """ Manages a sequential thinking session that records thoughts and allows for retrieval. Uses MCP to store and retrieve thoughts in a ChromaDB collection. """ def __init__(self, client: ClientSession = None, session_id: str = None, mcp_server_url: Optional[str] = None): """ Initialize a thinking session. Args: client: ClientSession instance for communicating with MCP server session_id: Optional session ID. If not provided, a new UUID will be generated. mcp_server_url: Optional MCP server URL. If not provided, it will be read from the environment variable MCP_SERVER_URL. """ if client: self.client = client else: server_url = mcp_server_url or os.getenv("MCP_SERVER_URL", "http://localhost:8000") if not server_url: raise ValueError("MCP_SERVER_URL environment variable not set and no client or URL provided.") self.client = ClientSession(server_url) self.session_id = session_id or str(uuid.uuid4()) def record_thought( self, thought: str, thought_number: int, total_thoughts: int, branch_id: str = "", branch_from_thought: int = 0, next_thought_needed: bool = False, ) -> Dict[str, Any]: """ Record a single thought in the thinking session. Args: thought: Content of the thought thought_number: Sequential number of this thought total_thoughts: Total anticipated number of thoughts in this sequence branch_id: Optional identifier for a branch within the session branch_from_thought: If creating a new branch, the parent thought number it originates from next_thought_needed: Flag indicating if a subsequent thought is expected Returns: Response from MCP server including the session ID """ response = self.client.mcp_chroma_dev_chroma_sequential_thinking( thought=thought, thought_number=thought_number, total_thoughts=total_thoughts, session_id=self.session_id, branch_id=branch_id, branch_from_thought=branch_from_thought, next_thought_needed=next_thought_needed, ) # Update session_id if it was auto-generated by the server if not self.session_id and "session_id" in response: self.session_id = response["session_id"] return response def find_similar_thoughts( self, query: str, n_results: int = 5, threshold: float = -1.0, include_branches: bool = True ) -> List[Dict[str, Any]]: """ Find thoughts semantically similar to the query. Args: query: Text to search for similar thoughts n_results: Maximum number of similar thoughts to return threshold: Similarity score threshold (0.0 to 1.0). Lower distance is more similar. -1.0 to use default. include_branches: Whether to include thoughts from branches in the search Returns: List of similar thoughts with metadata """ response = self.client.mcp_chroma_dev_chroma_find_similar_thoughts( query=query, session_id=self.session_id, n_results=n_results, threshold=threshold, include_branches=include_branches, ) return response.get("similar_thoughts", []) def get_session_summary(self, include_branches: bool = True) -> Dict[str, Any]: """ Get a summary of all thoughts in this thinking session. Args: include_branches: Whether to include thoughts from branches in the summary Returns: Dictionary containing session information and all thoughts """ response = self.client.mcp_chroma_dev_chroma_get_session_summary( session_id=self.session_id, include_branches=include_branches ) return response @classmethod def find_similar_sessions( cls, query: str, client: ClientSession = None, n_results: int = 5, threshold: float = -1.0 ) -> List[Dict[str, Any]]: """ Find thinking sessions similar to the query text. Args: query: Text to search for similar sessions client: ClientSession instance for communicating with MCP server n_results: Maximum number of similar sessions to return threshold: Similarity score threshold (0.0 to 1.0). Lower distance is more similar. -1.0 to use default. Returns: List of similar sessions with metadata """ client = client or ClientSession(os.getenv("MCP_SERVER_URL", "http://localhost:8000")) response = client.mcp_chroma_dev_chroma_find_similar_sessions( query=query, n_results=n_results, threshold=threshold ) return response.get("similar_sessions", [])