diff --git a/lab_notebook_intelligence/ai_service_manager.py b/lab_notebook_intelligence/ai_service_manager.py index 4007de1..4e57f70 100644 --- a/lab_notebook_intelligence/ai_service_manager.py +++ b/lab_notebook_intelligence/ai_service_manager.py @@ -137,6 +137,66 @@ def initialize(self): self.update_models_from_config() self.initialize_extensions() + # Create dynamic MCP config for user's working directory + self._create_dynamic_mcp_config() + + def _create_dynamic_mcp_config(self): + # Get the directory where JupyterLab was started (user's working directory) + user_root_dir = self._options.get("server_root_dir", os.getcwd()) + + log.info(f"Creating dynamic MCP config for directory: {user_root_dir}") + + # Create dynamic MCP config with filesystem servers + dynamic_mcp_config = { + "mcpServers": { + "filesystem-pwd": { + "command": "npx", + "args": [ + "-y", + "@modelcontextprotocol/server-filesystem", + user_root_dir, + ], + }, + "qbraid-web-search": { + "command": "uv", + "args": ["tool", "run", "web-browser-mcp-server"], + "env": { + "REQUEST_TIMEOUT": "60", + }, + }, + # add the MCP server for accessing docs.qbraid.com + # "qbraid-docs-search": {"url": "https://docs.qbraid.com/mcp"}, + "context7-search": { + "url": "https://mcp.context7.com/mcp", + "headers": {"CONTEXT7_API_KEY": os.getenv("CONTEXT7_API_KEY", "")}, + }, + } + } + + # Add qBraid environments MCP server + qbraid_envs_dir = os.path.expanduser("~/.qbraid/environments/") + log.info(f"qBraid environments directory: {qbraid_envs_dir}") + if os.path.exists(qbraid_envs_dir): + # Add filesystem access to environments directory + dynamic_mcp_config["mcpServers"]["qbraid-envs"] = { + "command": "npx", + "args": [ + "-y", + "@modelcontextprotocol/server-filesystem", + qbraid_envs_dir, + ], + } + log.info(f"Added qBraid environments MCP server") + + else: + log.info(f"qBraid environments directory not found: {qbraid_envs_dir}") + + # Save to user's MCP config (this will merge with existing config) + self.nbi_config.user_mcp = dynamic_mcp_config + self.nbi_config.save() + self.nbi_config.load() + self.update_mcp_servers() + def update_models_from_config(self): using_github_copilot_service = self.nbi_config.using_github_copilot_service if using_github_copilot_service: @@ -253,6 +313,7 @@ def register_completion_context_provider(self, provider: CompletionContextProvid log.error(f"Completion Context Provider ID '{provider.id}' is already in use!") return self.completion_context_providers[provider.id] = provider + log.info(f"Registered completion context provider: {provider.id}") def register_telemetry_listener(self, listener: TelemetryListener) -> None: if listener.name in self.telemetry_listeners: @@ -439,7 +500,8 @@ async def get_completion_context(self, request: ContextRequest) -> CompletionCon if cancel_token.is_cancel_requested: return context - + log.debug(f"Allowed context providers: {allowed_context_providers}") + log.debug(f"Available providers: {list(self.completion_context_providers.keys())}") for provider in self.completion_context_providers: if cancel_token.is_cancel_requested: return context @@ -450,7 +512,13 @@ async def get_completion_context(self, request: ContextRequest) -> CompletionCon ): continue try: - provider_context = provider.handle_completion_context_request(request) + # Try async version first, fallback to sync + if hasattr(provider, "async_handle_completion_context_request"): + provider_context = await provider.async_handle_completion_context_request( + request + ) + else: + provider_context = provider.handle_completion_context_request(request) if provider_context.items: context.items += provider_context.items except Exception as e: diff --git a/lab_notebook_intelligence/base_chat_participant.py b/lab_notebook_intelligence/base_chat_participant.py index fdeeaea..92537de 100644 --- a/lab_notebook_intelligence/base_chat_participant.py +++ b/lab_notebook_intelligence/base_chat_participant.py @@ -4,6 +4,7 @@ import json import logging import os +import re from typing import Union from lab_notebook_intelligence.api import ( @@ -383,6 +384,177 @@ async def handle_tool_call( return {"result": "Code cell added to notebook"} +class AutoLibraryDocsTool(Tool): + @property + def name(self) -> str: + return "auto_library_docs" + + @property + def title(self) -> str: + return "Auto Library Documentation" + + @property + def tags(self) -> list[str]: + return ["internal-tool", "documentation"] + + @property + def description(self) -> str: + return "Automatically fetch library documentation when library-related queries are detected" + + @property + def schema(self) -> dict: + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "strict": False, + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The user query that needs library documentation", + }, + "libraries": { + "type": "array", + "items": {"type": "string"}, + "description": "List of specific libraries detected in the query", + }, + "search_terms": { + "type": "array", + "items": {"type": "string"}, + "description": "Broader search terms for finding relevant library documentation", + }, + }, + "required": ["query"], + "additionalProperties": False, + }, + }, + } + + def pre_invoke( + self, request: ChatRequest, tool_args: dict + ) -> Union[ToolPreInvokeResponse, None]: + # This tool runs automatically, no user confirmation needed + return None + + async def handle_tool_call( + self, + request: ChatRequest, + response: ChatResponse, + tool_context: dict, + tool_args: dict, + ) -> str: + query = tool_args.get("query", "") + libraries = tool_args.get("libraries", []) + search_terms = tool_args.get("search_terms", []) + + # Combine libraries and search terms for comprehensive search + all_search_items = libraries + search_terms + + if not all_search_items: + return "No libraries or search terms provided" + + documentation_results = [] + + # Get MCP servers from host + host = request.host + try: + mcp_servers = host.get_mcp_servers() + except Exception as e: + return f"Could not get MCP servers: {e}" + + # Find servers that have the required tools + resolve_server = None + docs_server = None + for server in mcp_servers: + server_tools = [tool.name for tool in server.get_tools()] + if "resolve-library-id" in server_tools: + resolve_server = server + if "get-library-docs" in server_tools: + docs_server = server + + if not resolve_server or not docs_server: + return "Required MCP tools (resolve-library-id, get-library-docs) not found" + + for search_item in all_search_items: + try: + log.info(f"Attempting to resolve library: {search_item}") + + # Step 1: Resolve library ID using the search item + resolve_result = await resolve_server.call_tool( + "resolve-library-id", {"libraryName": search_item} + ) + + if not resolve_result or not hasattr(resolve_result, "content"): + log.info(f"No resolve result for '{search_item}'") + continue + + # Extract library ID from result + library_id = None + if isinstance(resolve_result.content, list) and len(resolve_result.content) > 0: + content = resolve_result.content[0] + if hasattr(content, "text"): + # Parse the library ID from the response + import json + + try: + parsed = json.loads(content.text) + library_id = ( + parsed.get("library_id") + or parsed.get("id") + or parsed.get("libraryId") + ) + log.info(f"Resolved '{search_item}' to library_id: {library_id}") + except: + # If not JSON, assume the text is the library ID + library_id = content.text.strip() + log.info(f"Using text as library_id for '{search_item}': {library_id}") + + if not library_id: + log.info(f"Could not extract library ID for '{search_item}'") + continue + + # Step 2: Get library documentation + log.info(f"Fetching docs for library_id: {library_id}") + docs_result = await docs_server.call_tool( + "get-library-docs", {"context7CompatibleLibraryID": library_id} + ) + + if ( + docs_result + and hasattr(docs_result, "content") + and isinstance(docs_result.content, list) + ): + docs_content = [] + for content in docs_result.content: + if hasattr(content, "text"): + docs_content.append(content.text) + + if docs_content: + log.info(f"Successfully fetched docs for '{search_item}'") + documentation_results.append( + f"## Documentation for '{search_item}'\n\n" + "\n\n".join(docs_content) + ) + else: + log.info(f"No documentation content found for '{search_item}'") + else: + log.info(f"Could not retrieve documentation for '{search_item}'") + + except Exception as e: + log.error(f"Error fetching documentation for '{search_item}': {str(e)}") + + if documentation_results: + # Add the documentation to the response context + docs_text = "\n\n".join(documentation_results) + response.stream(MarkdownData(docs_text)) + return f"Successfully retrieved documentation for some search terms" + else: + log.info("No documentation could be retrieved for any search terms") + return "No documentation could be retrieved" + + class BaseChatParticipant(ChatParticipant): def __init__(self): super().__init__() @@ -460,10 +632,12 @@ async def generate_code_cell(self, request: ChatRequest) -> str: 0, { "role": "system", - "content": f"You are an assistant that creates Python code which will be used in a Jupyter notebook. Generate only Python code and some comments for the code. You should return the code directly, without wrapping it inside ```.", + "content": f"You are an assistant that creates correct and executable Python code which will be used in a Jupyter notebook. Whenever you are using multiple libraries, ensure that the generated code is compatible and does not use deprecated or non-existent methods. Generate only Python code and some comments for the code. You should return the code directly, without wrapping it inside ```.", }, ) - messages.append({"role": "user", "content": f"Generate code for: {request.prompt}"}) + messages.append( + {"role": "user", "content": f"Generate clean Python code for: {request.prompt}"} + ) generated = chat_model.completions(messages) code = generated["choices"][0]["message"]["content"] @@ -491,10 +665,44 @@ async def generate_markdown_for_code(self, request: ChatRequest, code: str) -> s return extract_llm_generated_code(markdown) + async def generate_title_for_notebook( + self, request: ChatRequest, code: str, markdown: str + ) -> str: + chat_model = request.host.chat_model + messages = [ + { + "role": "system", + "content": f"You are an assistant that generates descriptive titles for Jupyter notebooks based on the provided code and markdown. The title should be in snake-case, concise, relevant, and accurately reflect the content of the notebook. Avoid using special characters, use '_' instead of spaces and keep it under 50 characters.", + }, + { + "role": "user", + "content": f"Generate a descriptive title for a Jupyter notebook that contains the following markdown and code:\n\nMarkdown:\n{markdown}\n\nCode:\n{code}\n\n", + }, + ] + generated = chat_model.completions(messages) + + # if title is too long, truncate it + title = generated["choices"][0]["message"]["content"].strip().strip('"').strip("'")[:50] + + # if this notebook already exists, append a number to the title + server_root_dir = request.host.nbi_config.server_root_dir + file_path = os.path.join(server_root_dir, f"{title}.ipynb") + counter = 1 + while os.path.exists(file_path): + title = f"{title}_{counter}" + file_path = os.path.join(server_root_dir, f"{title}.ipynb") + counter += 1 + + return title + async def handle_chat_request( self, request: ChatRequest, response: ChatResponse, options: dict = {} ) -> None: self._current_chat_request = request + + # Auto-detect libraries in the query and fetch documentation + await self._auto_fetch_library_docs(request, response) + if request.chat_mode.id == "ask": return await self.handle_ask_mode_chat_request(request, response, options) elif request.chat_mode.id == "agent": @@ -527,6 +735,90 @@ async def handle_chat_request( await self.handle_chat_request_with_tools(request, response, options) + async def _auto_fetch_library_docs(self, request: ChatRequest, response: ChatResponse) -> None: + """Use LLM to intelligently detect if library documentation is needed""" + if not request.prompt: + log.info("No prompt provided, skipping library docs fetch") + return + + log.info(f"Checking if library docs needed for prompt: {request.prompt[:100]}") + + # Check if we have the required MCP tools available + host = request.host + # Use the get_mcp_servers() method from the host (AIServiceManager) + try: + mcp_servers = host.get_mcp_servers() + except Exception as e: + log.info(f"Could not get MCP servers: {e}") + return + + has_resolve_tool = False + has_docs_tool = False + for server in mcp_servers: + server_tools = [tool.name for tool in server.get_tools()] + if "resolve-library-id" in server_tools: + has_resolve_tool = True + if "get-library-docs" in server_tools: + has_docs_tool = True + + log.info(f"Has resolve tool: {has_resolve_tool}, Has docs tool: {has_docs_tool}") + + if not (has_resolve_tool and has_docs_tool): + log.info("Required MCP tools not available, skipping library docs fetch") + return + + # Use LLM to determine if library documentation is needed + try: + chat_model = request.host.chat_model + detection_result = await self._llm_detect_library_need(chat_model, request.prompt) + + log.info(f"Detection result: {detection_result}") + if detection_result and detection_result.get("needs_docs", False): + log.info("Library docs needed, fetching...") + # Create and execute the auto library docs tool + auto_docs_tool = AutoLibraryDocsTool() + await auto_docs_tool.handle_tool_call( + request, + response, + {}, + { + "query": request.prompt, + "libraries": detection_result.get("libraries", []), + "search_terms": detection_result.get("search_terms", []), + }, + ) + else: + log.info("Library docs not needed for this query") + except Exception as e: + log.warning(f"Failed to auto-fetch library documentation: {e}") + + async def _llm_detect_library_need(self, chat_model, query: str) -> dict: + """Use LLM to detect if a query needs library documentation""" + detection_prompt = Prompts.library_detection_prompt(query) + + try: + messages = [{"role": "user", "content": detection_prompt}] + response = chat_model.completions(messages) + + if response and "choices" in response and len(response["choices"]) > 0: + content = response["choices"][0]["message"]["content"].strip() + + # Try to parse JSON response + import json + + try: + return json.loads(content) + except json.JSONDecodeError: + # If JSON parsing fails, try to extract JSON from the response + json_match = re.search(r"\{.*\}", content, re.DOTALL) + if json_match: + return json.loads(json_match.group()) + + except Exception as e: + log.warning(f"LLM library detection failed: {e}") + + return {"needs_docs": False, "libraries": [], "search_terms": []} + async def handle_ask_mode_chat_request( self, request: ChatRequest, response: ChatResponse, options: dict = {} ) -> None: @@ -550,7 +842,16 @@ async def handle_ask_mode_chat_request( {"code": code, "path": file_path}, ) - response.stream(MarkdownData(f"Notebook '{file_path}' created and opened successfully")) + # get a descriptive name for the notebook + title = await self.generate_title_for_notebook(request, code, markdown) + ui_cmd_response = await response.run_ui_command( + "lab-notebook-intelligence:rename-notebook", {"newName": title} + ) + new_file_path = ui_cmd_response.get("newPath", file_path) + + response.stream( + MarkdownData(f"Notebook '{new_file_path}' created and opened successfully") + ) response.finish() return elif request.command == "newPythonFile": @@ -561,7 +862,7 @@ async def handle_ask_mode_chat_request( 0, { "role": "system", - "content": f"You are an assistant that creates Python code. You should return the code directly, without wrapping it inside ```.", + "content": f"You are an assistant that creates correct and executable Python code. You should return the code directly, without wrapping it inside ```.", }, ) messages.append({"role": "user", "content": f"Generate code for: {request.prompt}"}) @@ -614,5 +915,7 @@ def get_tool_by_name(name: str) -> Tool: return AddMarkdownCellToNotebookTool() elif name == "add_code_cell_to_notebook": return AddCodeCellTool() + elif name == "auto_library_docs": + return AutoLibraryDocsTool() return None diff --git a/lab_notebook_intelligence/extension.py b/lab_notebook_intelligence/extension.py index 77911bd..30362cd 100644 --- a/lab_notebook_intelligence/extension.py +++ b/lab_notebook_intelligence/extension.py @@ -211,108 +211,6 @@ def post(self): return -class CreateDynamicMCPConfigHandler(APIHandler): - @tornado.web.authenticated - def post(self): - try: - # Get the directory where JupyterLab was started (user's working directory) - user_root_dir = NotebookIntelligence.root_dir - - print(f"Creating dynamic MCP config for directory: {user_root_dir}") - - # Create dynamic MCP config with filesystem servers - dynamic_mcp_config = { - "mcpServers": { - "filesystem-pwd": { - "command": "npx", - "args": [ - "-y", - "@modelcontextprotocol/server-filesystem", - user_root_dir, - ], - }, - "qbraid-web-search": { - "command": "uv", - "args": ["tool", "run", "web-browser-mcp-server"], - "env": { - "REQUEST_TIMEOUT": "60", - }, - }, - # add the MCP server for accessing docs.qbraid.com - "qbraid-docs-search": {"url": "https://docs.qbraid.com/mcp"}, - } - } - - # Add qBraid environments MCP server - qbraid_envs_dir = os.path.expanduser("~/.qbraid/environments/") - print(f"qBraid environments directory: {qbraid_envs_dir}") - if os.path.exists(qbraid_envs_dir): - # Add filesystem access to environments directory - dynamic_mcp_config["mcpServers"]["qbraid-envs"] = { - "command": "npx", - "args": [ - "-y", - "@modelcontextprotocol/server-filesystem", - qbraid_envs_dir, - ], - } - print(f"Added qBraid environments MCP server") - - # TODO: Uncomment and fix the code below to add individual Python execution servers for each qBraid environment - - # # Discover individual environments and add Python execution servers - # try: - # env_count = 0 - # for env_name in os.listdir(qbraid_envs_dir): - # env_path = os.path.join(qbraid_envs_dir, env_name) - # python_executable = os.path.join(env_path, "bin", "python") - - # # Check if this is a valid environment with Python - # if (os.path.isdir(env_path) and - # os.path.exists(python_executable) and - # not env_name.startswith('.')): - - # # Add Python execution server for this environment - # server_name = f"python-{env_name}" - # dynamic_mcp_config["mcpServers"][server_name] = { - # "command": "npx", - # "args": [ - # "-y", - # "@modelcontextprotocol/python-mcp" - # ], - # "env": { - # "PYTHON_EXECUTABLE": python_executable - # } - # } - # env_count += 1 - # print(f"Added Python execution server for environment: {env_name}") - - # print(f"Discovered and added {env_count} qBraid environment Python servers") - # except Exception as e: - # print(f"Error discovering qBraid environments: {e}") - - else: - print(f"qBraid environments directory not found: {qbraid_envs_dir}") - - # Save to user's MCP config (this will merge with existing config) - ai_service_manager.nbi_config.user_mcp = dynamic_mcp_config - ai_service_manager.nbi_config.save() - ai_service_manager.nbi_config.load() - ai_service_manager.update_mcp_servers() - - self.finish( - json.dumps( - { - "status": "ok", - "message": f"Dynamic MCP config created for directory: {user_root_dir}", - } - ) - ) - except Exception as e: - self.finish(json.dumps({"status": "error", "message": str(e)})) - return - - class EmitTelemetryEventHandler(APIHandler): @tornado.web.authenticated def post(self): @@ -960,9 +858,6 @@ def _setup_handlers(self, web_app): route_pattern_mcp_config_file = url_path_join( base_url, "lab-notebook-intelligence", "mcp-config-file" ) - route_pattern_create_dynamic_mcp_config = url_path_join( - base_url, "lab-notebook-intelligence", "create-dynamic-mcp-config" - ) route_pattern_emit_telemetry_event = url_path_join( base_url, "lab-notebook-intelligence", "emit-telemetry-event" ) @@ -983,7 +878,6 @@ def _setup_handlers(self, web_app): (route_pattern_update_provider_models, UpdateProviderModelsHandler), (route_pattern_reload_mcp_servers, ReloadMCPServersHandler), (route_pattern_mcp_config_file, MCPConfigFileHandler), - (route_pattern_create_dynamic_mcp_config, CreateDynamicMCPConfigHandler), (route_pattern_emit_telemetry_event, EmitTelemetryEventHandler), (route_pattern_github_login_status, GetGitHubLoginStatusHandler), (route_pattern_github_login, PostGitHubLoginHandler), diff --git a/lab_notebook_intelligence/mcp_manager.py b/lab_notebook_intelligence/mcp_manager.py index be713e1..7d4b3f3 100644 --- a/lab_notebook_intelligence/mcp_manager.py +++ b/lab_notebook_intelligence/mcp_manager.py @@ -17,7 +17,6 @@ ChatCommand, ChatRequest, ChatResponse, - HTMLFrameData, ImageData, MarkdownData, MCPServer, @@ -32,7 +31,7 @@ MCP_ICON_SRC = "iVBORw0KGgoAAAANSUhEUgAAAMgAAADICAIAAAAiOjnJAAAPBUlEQVR4nOydf2wT5f/AW7pZGLOjE7K5DAfWIYMWM7rpWJTZkCxEh4hdcGKahYhkYRojGPQfUnCJMRpDlpD5hyEknZlWY2AL2UaiDubYLFkDzsE2CBlBrc5ldGm6VLjdnm++6Sf77DO758fdPb279v36kzz3ft73vhfPdffjfRkIIQMAKM0ytRMAUhMQC+ACiAVwAcQCuABiAVwAsQAugFgAF0AsgAsgFsAFEAvgAogFcAHEArgAYgFcALEALoBYABdALIALIBbABRAL4AKIBXABxAK4kKF2Anrlt99+6+vru3r16s2bN+/cuTM1NRWJRGZnZ1esWGGxWPLy8mw22+bNm8vLy7dt27Zy5Uq18002RnhLh4mhoaGvvvqqo6Pjxo0blJuYzebKykq32/3qq6+uXr2ac4KaAQEUiKL4zTffVFRUyCm12Wyur6+/du2a2nuTDEAsMu3t7Xa7Xan/ySaTae/evbdu3VJ7t/gCYuG4e/duTU2NUkotJCsry+v1CoKg9i7yAsRaEr/fb7VaeVg1zzPPPDM+Pq72jnIBxEqAKIoffPABV6XmWbVqVU9Pj9p7rDwg1mIEQairq0uOVXFMJlNbW5va+60wINb/cP/+/d27dyfTqlR1C8T6L2pZNe/W2bNn1a6BYsAF0v/w4MGDvXv3tre3q5hDdnb25cuXt2zZomIOSgFiGeRYZbVan3/++a1bt27YsCE3NzcjI0MQhD/++OP69es///xzIBAQBIEpYHFxcTAYfPjhh1kz0RxqL5nqI+EMaDQa3W53V1cX/kJUOBw+derUxo0bmYIfOnQoiXvPi3QXS4JVLpdraGiIfgpBEHw+X35+Pr21vb29PHc6GaS1WKxWmc3mU6dOSZtrenra7XZTTuRwOERRVHp3k0r6isVqldVq7evrkznpsWPHKKfz+/0K7ag6pKlYEqwKBoOKTH38+HGaGe12uyLTqUU6iqWiVXEaGhpo5tX1rZ60E0t1qxBCsVhs06ZNxKk9Ho+y8yaT9BJLC1bFGRgYMBqN+NktFkssFuMxexJII7G0Y1Wc2tpaYg4XLlzglwBX0uUtHdZr61ar9fvvv9+6dSu/lN5//33imJ6eHn4JcCUt3tLhbdXc3NyPP/7Y29sbDofXr1//8ssvP/7448StysrKHA7Hr7/+ihkzMDBAmYPmUHvJ5A7vM+CtW7fKy8sXRjCZTI2Njffv3ydu6/V6icnI23vVSHGxeFs1NjZWUFCQMFRdXR1x8wsXLhBTmpyclFcDdUhlsVS0Ks63336LjzA5OUnM6pdffpFdCRVIWbFUt8pgMOzatYsYZ/ny5fggOr1MmppiacEqg8FQWFhIDEWMc/78eXnFUIcUvNzA+2/AmzdvulyuUChEHDk3N6fIGD2SamJpxyqDwUDzkHEkEsEPWLFiBWVu2kLtJVNJNHIGnOe7777DB5yYmCAGuXr1quzCqEDqXCDV1FplMBh27979yiuv4McMDw8T42RmZg4ODobD4QcPHsQfNszNzS0oKKB/JFUd1DZbGbS2VlVVVUWjUWJY4gVSDBaLZfv27UePHu3s7NTgvepUEEunViGESktLJUm1GIvF8vrrr2vqwoTuxdKvVcFgUJJFODZt2nTmzBktNLHRt1j6tQoh5PF4JMlDprS09NKlS1KLqgw6FkvXVgUCAZPJJEkbWt54443p6Wmp1ZWLXsXStVWCIDgcDkm2sFFcXKzWrUZdiqVrqxBCjY2NkjyRgsViUaXXiP7E0rtVH3/8sSRDpKNKjySdiaV3qz755BNJbsgl+W7pSSywSg4mk+ncuXOSCi8F3YiV5lZlZ2c/9thjhYWFFotFThCmdiZy0IdY6WmV0+lsamq6dOnS1NTUwmhTU1M9PT1er/epp55ijblhw4ZIJMJSe4noQKx0s8poNHo8HsqlJRAI1NbWMl0SS07/La2LlW5WlZeXS/gmSiAQoL8wlpz+W5oWK92sOnz4sOTbfLFY7ODBg5QTlZaW8u6/pV2x0soqo9HY0tIiqU7/w4cffkg5I+/+WxoVK92sOn36tKQ6JYDSLYfDodSMCdGiWGCVTCjPifIbFGLQnFhglXwo+2/V19crPvU82hILrFKKvr4+Yv+tnJwcmgYT0tCQWGCVstD03/rhhx84za4VscAqxbl8+TIxk2PHjnGaXRNigVUY4p+j3rFjR35+fkFBQXV1Nf3zVU888QQ+merqavpMmFBfLLAKQzQafeGFF/4dp7a2lubnEfFznvn5+fTJMKGyWGAVhmg0WlVVtVS0t956ixihs7OTmBWn5+LVFAuswoC3Kh6Q+Cl8mve2b9y4QZ8VPaqJBVZhIFoV5+TJk/g4oihmZmbig3C6TKqOWGAVBkqrDAbDkSNHiNGIH+Ln9P60CmKBVRjorTIYDCdOnCAGfOSRR/BBUkQssAoDk1U0ToiiSHwGMBVOhWAVBlarKioqiDHv3r1LjKP7H+9gFQZWq4qKisbHx4lhz58/TwzF6RH4JIkFVmHgZBVC6OjRo/hQ+r5AClZh4GcVQqi4uBgfbefOnfSpMsFdLLAKA1er+vv7iQG9Xi99tkzwFQuswsDVKoRQwpuMi+DXRoujWGAVBt5W9fb2EmPm5OTw6/3HSyywCgNvq6LR6JNPPkkMe+DAAfqYrHARC6zCwNsqhND+/ftpIvf39zOFZUJ5sQRB2LNnD33hWK0aHx8HqzBQtvguKytjCsuK8mIx9S5ntWpqaqqoqIg+Pli1FLy//aSwWNPT01lZWZT7xmqVIAgul4v+wIBVS8F7uVJerO7ubsp9k/CN+BMnTtAfGLBqKUwmE9dfV3EUFuvMmTM0+ybBqmvXrtE36wGrMLz99ttMwaWhwoolwSpBEMrKyigLB1Zh2LhxI1NxJKOwWJFIJCcnB7NjEqxCCPl8PsrCgVUYLBYLp4dk/o3yfxVi2k1Ls0oQBJvNRlM4sAqDyWTq7Oxkii8HLtex9u3b9+8dy8vLk2AVQujs2bM0hausrASrliJ12nH7fD6n0xnvS5GXl3fo0KFQKCQt1I4dO4iFKywsnJiYoI8JVvGG79MNsVgsHA7LiRAKhWj+GOzq6qKPCVYlAfVfscfT0tJCrJ3b7aYPCFYlB62LVVNTQ6zdyMgIZTSwKmloWixRFFetWoUvH32/FLAqmWharJGREWIFfT4fTSiwKsloWqxz584Ri0jzxyBYlXw0LdZnn32GL6LNZiMG+fTTT+mPClilFJoW68iRI/g6vvjii/gIXV1dxB6v84BVCqJpsYj9yg8ePIiP4HQ6KY8KWKUsmharvr4eX83GxkbM5jSdC+KAVYqjabGILwU0NDRgNqd5BQqs4oSmxXrnnXfwNa2trcVsPjw8TDwqYBUnNC1WU1MTvqxbtmzBbC6KYn5+PmZzsIofmhbryy+/xFc2MzNzZmYGE+HkyZNglSpoWqwrV64Q69vd3Y2JIIpiXV1dwqMCVnFF02LNzMwQm/7u378fH0QQhObm5oXnRKfTyfQNGbBKApoWCyFUUVGBr3J2dva9e/eIcURRHB0dvXLlCuvzhmCVNLQuFk3R+X1pCKySjNbFCgaDxHJnZWXdvn1b8anBKjloXSyEEM3HQquqqpRt9QRWyUQHYmEuGSzk3XffVWpGsEo+OhDr3r17+Jdg51Gko2YkEgGr5KMDsZiOxOHDh+WcE0Oh0NNPPw1WyUcfYk1PT+Nvzixk+/btxO+tJaS9vZ1+FrAKjz7EYmrfEP878b333qN/hXVoaIipCyFYRUQ3YiGEdu3axXTss7KyPB5PR0fHUq/eh0KhL774wuVy0T9lClZRYvx/uXTC33//7XQ6f//9d9YNMzMzS0pKbDbbmjVrMjIyYrHYn3/+OTo6eufOHQlpFBUVXbx4cd26dfSbHD9+nL5rnMlkam1tfe211yTkpiHUNpuN/v7+5cuXq1guWKso0ZlYCCG/30/f2k9ZwCp69CcWQqitrS35boFVTOhSrPi6lcxzYnFxMVjFhF7FQggNDAwUFhby1Ok/VFdXszZjSnOr9C0WQmhiYoLYjkYOZrP5o48+EkWRKSuwSvdixWlra1uzZo3iVj377LPDw8OsyYBVcVJBrPjzCF6v12q1KqKU3W73+/2sCxVYtZAUEStONBptbm622+3SfDIajTU1NR0dHRKUAqsWkVJizRMMBpuamiorK81mM/EYW63WPXv2tLS0/PXXX5JnBKsWoadbOhL4559/RkdHr1+/HgqFJiYmZmZmYrHYypUrLRbLo48+un79+s2bN69bt27ZsmVyZknHOzZE1DZb98BalRAQSxZg1VKAWNIBqzCAWBIBq/CAWFIAq4iAWMyAVTSAWGyAVZSAWAyAVfSAWLSAVUyAWFSAVayAWGTAKgmAWAQoW5KAVYtI8ZvQMvnpp59cLpcoijSD0+XuMh0g1pLMzs7a7faxsTGawWDVImQ9LpLanD59GqySDKxYS1JSUjI6OkocBlYlBFasxIyMjIBVcgCxEnPx4kXiGLAKA4iVGOLnqMEqPCBWYiYnJ/ED3G43WIUBxErM3NwcfsDq1auTlYsuAbESk5ubix/g8/kGBgaSlY7+ALESY7PZ8AOi0ejOnTvBraUAsRLz3HPPEcdEIhFwayngAmli5ubm1q5dGwqFiCMtFkt3d/e2bduSkpdugBUrMcuWLXvzzTdpRsK6lRBYsZYkHA7bbLZwOEwzGNatRcCKtSRWq7W5uZlyMKxbiwCxcHg8ngMHDlAOBrcWAqdCArOzsx6P5+uvv6YcD+fEOLBiEcjIyGhtbU34KfyEwLoVB8QiA25JAMSiAtxiBX5jMSDh91ZfX5/D4eCclxYBsdhgdcvpdA4ODnJOSovAqZAN1nNiMBgEsQAqWN0KBAKcM9IiIJYUmNwSBIF/RpoDxJIIvVslJSVJyUhbwI93WRB/yxcUFNy+fVvdr8KqAqxYssCvW0aj8fPPP09Dq0AsBYi71dDQsOjfs7OzW1tbX3rpJZXyUhk4FSrG4OCg3+8fGxt76KGHysvL9+3bt3btWrWTUg0QC+ACnAoBLoBYABdALIALIBbABRAL4AKIBXABxAK4AGIBXACxAC6AWAAXQCyACyAWwAUQC+ACiAVwAcQCuABiAVz4vwAAAP//b8cbMGXTzMEAAAAASUVORK5CYII=" MCP_ICON_URL = f"data:image/png;base64,{MCP_ICON_SRC}" MCP_TOOL_TIMEOUT = 60 -WHITELISTED_MCP_TOOLS = {"SearchQBraid"} +WHITELISTED_MCP_TOOLS = {"resolve-library-id", "get-library-docs"} class MCPTool(Tool): @@ -264,6 +263,12 @@ def servers(self) -> list[MCPServer]: async def handle_chat_request( self, request: ChatRequest, response: ChatResponse, options: dict = {} ) -> None: + self._current_chat_request = request + + # Auto-detect libraries in the query and fetch documentation + await self._auto_fetch_library_docs(request, response) + log.info("Completed auto-fetch library docs check") + response.stream(ProgressData("Thinking...")) if request.command == "info": @@ -301,6 +306,13 @@ def update_mcp_servers(self, mcp_config): participant_servers = self.create_servers(server_names, servers_config) if len(participant_servers) > 0: + # Auto-add the library docs tool if context7 tools are available + if ( + self._has_context7_tools(participant_servers) + and "auto_library_docs" not in nbi_tools + ): + nbi_tools.append("auto_library_docs") + self._mcp_participants.append( MCPChatParticipant( f"mcp-{participant_id}", @@ -327,6 +339,11 @@ def update_mcp_servers(self, mcp_config): unused_servers = self.create_servers(unused_server_names, servers_config) mcp_participant_config = participants_config.get("mcp", {}) nbi_tools = mcp_participant_config.get("nbiTools", []) + + # Auto-add the library docs tool if context7 tools are available + if self._has_context7_tools(unused_servers) and "auto_library_docs" not in nbi_tools: + nbi_tools.append("auto_library_docs") + self._mcp_participants.append( MCPChatParticipant("mcp", "MCP", unused_servers, nbi_tools) ) @@ -335,6 +352,20 @@ def update_mcp_servers(self, mcp_config): thread = threading.Thread(target=self.init_tool_lists, args=()) thread.start() + def _has_context7_tools(self, servers: list[MCPServer]) -> bool: + """Check if any server has both resolve-library-id and get-library-docs tools""" + has_resolve = False + has_docs = False + + for server in servers: + server_tools = [tool.name for tool in server.get_tools()] + if "resolve-library-id" in server_tools: + has_resolve = True + if "get-library-docs" in server_tools: + has_docs = True + + return has_resolve and has_docs + def create_servers(self, server_names: list[str], servers_config: dict): servers = [] for server_name in server_names: diff --git a/lab_notebook_intelligence/prompts.py b/lab_notebook_intelligence/prompts.py index 4741a25..dd5ead7 100644 --- a/lab_notebook_intelligence/prompts.py +++ b/lab_notebook_intelligence/prompts.py @@ -5,7 +5,7 @@ OS_TYPE = "Linux" CHAT_SYSTEM_PROMPT = """ -You are an AI programming assistant. +You are an AI programming assistant for qBraid jupyter lab users. When asked for your name, you must respond with "{AI_ASSISTANT_NAME}". Follow the user's requirements carefully & to the letter. Follow Microsoft content policies. @@ -27,6 +27,7 @@ * Ask how to do something in the terminal * Explain what just happened in the terminal You use the {MODEL_NAME} AI model provided by {MODEL_PROVIDER}. +If there is any mention of "qbraid", in any capitalization, or related terms you should use the external MCP tool called "context7-search". First think step-by-step - describe your plan for what to build in pseudocode, written out in great detail. Then output the code in a single code block. This code block should not contain line numbers (line numbers are not necessary for the code to be understood, they are in format number: at beginning of lines). Minimize any other prose. @@ -39,6 +40,39 @@ You can only give one reply for each conversation turn. """ +CONTEXT_DOCS_IDENTIFICATION_PROMPT = """ +Analyze the following user query and determine if it would benefit from Python library documentation or platform-specific documentation. + +User Query: "{USER_QUERY}" + +Respond with a JSON object containing: +1. "needs_docs": boolean - true if the query is asking about Python libraries, packages, platforms, or programming concepts that would benefit from documentation +2. "libraries": array of strings - **ONLY** specific Python library/package names mentioned (e.g., ["pandas", "numpy", "qiskit", "braket", "qbraid"]) +3. "search_terms": array of strings - You can include strings which are required for this search e.g. "runtime", "quantum computing", "machine learning", etc. + +**Examples:** + +Query: "How can I submit a braket circuit to a qiskit backend using qbraid runtime?" +Response: {{"needs_docs": true, "libraries": ["braket", "qiskit", "qbraid"], "search_terms": ["braket", "qiskit", "qbraid"]}} + +Query: "What do you know about qbraid runtime?" +Response: {{"needs_docs": true, "libraries": ["qbraid"], "search_terms": ["qbraid"]}} + +Query: "How to use numpy arrays?" +Response: {{"needs_docs": true, "libraries": ["numpy"], "search_terms": ["numpy"]}} + +Query: "I'm getting an error with pandas DataFrame" +Response: {{"needs_docs": true, "libraries": ["pandas"], "search_terms": ["pandas"]}} + +Query: "How do I plot data?" +Response: {{"needs_docs": true, "libraries": [], "search_terms": ["matplotlib", "plotly", "seaborn"]}} + +Query: "What's the weather today?" +Response: {{"needs_docs": false, "libraries": [], "search_terms": []}} + +Only respond with the JSON object, no other text: +""" + class Prompts: @staticmethod @@ -60,3 +94,7 @@ def github_copilot_chat_prompt(model_provider: str, model_name: str) -> str: MODEL_NAME=model_name, MODEL_PROVIDER=model_provider, ) + + @staticmethod + def library_detection_prompt(user_query: str) -> str: + return CONTEXT_DOCS_IDENTIFICATION_PROMPT.format(USER_QUERY=user_query) diff --git a/src/api.ts b/src/api.ts index 84a91fb..5e05e6b 100644 --- a/src/api.ts +++ b/src/api.ts @@ -288,22 +288,6 @@ export class NBIAPI { }); } - static async createDynamicMCPConfig(): Promise { - return new Promise((resolve, reject) => { - requestAPI('create-dynamic-mcp-config', { - method: 'POST', - body: JSON.stringify({}) - }) - .then(async data => { - resolve(data); - }) - .catch(reason => { - console.error(`Failed to create dynamic MCP config.\n${reason}`); - reject(reason); - }); - }); - } - static async chatRequest( messageId: string, chatId: string, diff --git a/src/chat-sidebar.tsx b/src/chat-sidebar.tsx index ff96ca2..6f8e6c5 100644 --- a/src/chat-sidebar.tsx +++ b/src/chat-sidebar.tsx @@ -761,7 +761,6 @@ function SidebarComponent(props: any) { const uniqueFiles = Array.from(new Set(finalFiles)); setSelectedFiles(uniqueFiles); setShowFileBrowser(false); - console.log(`Selected files: ${filePaths}`); }; const handleContextCancel = () => { diff --git a/src/index.ts b/src/index.ts index b71f0f9..5cf723f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -135,6 +135,8 @@ namespace CommandIDs { const DOCUMENT_WATCH_INTERVAL = 1000; const MAX_TOKENS = 4096; +const STREAM_TOKEN_DELAY = 75; // milliseconds +const STREAM_TOKEN_LEN = 20; // characters const githubCopilotIcon = new LabIcon({ name: 'lab-notebook-intelligence:github-copilot-icon', svgstr: copilotSvgstr @@ -723,15 +725,6 @@ const plugin: JupyterFrontEndPlugin = { await NBIAPI.initialize(); - // Create dynamic MCP config after API is initialized - console.log('Creating dynamic MCP config...'); - try { - await NBIAPI.createDynamicMCPConfig(); - console.log('Dynamic MCP config created successfully!'); - } catch (error) { - console.error('Failed to create dynamic MCP config:', error); - } - let openPopover: InlinePromptWidget | null = null; let mcpConfigEditor: MCPConfigEditor | null = null; @@ -983,7 +976,7 @@ const plugin: JupyterFrontEndPlugin = { try { await app.serviceManager.contents.rename(oldPath, newPath); - return 'Successfully renamed notebook'; + return { newPath: newPath }; } catch (error) { return `Failed to rename notebook: ${error}`; } @@ -1051,12 +1044,25 @@ const plugin: JupyterFrontEndPlugin = { const newCellIndex = isNewEmptyNotebook(model) ? 0 : model.cells.length - 1; + model.insertCell(newCellIndex, { cell_type: cellType, metadata: { trusted: true }, - source + source: '' }); + const cell = currentWidget.content.widgets[newCellIndex]; + let currentText = ''; + (async () => { + for (let i = 0; i < source.length; i += STREAM_TOKEN_LEN) { + currentText += source.slice(i, i + STREAM_TOKEN_LEN); + cell.model.sharedModel.source = currentText; + // Wait for a short delay to simulate streaming + // eslint-disable-next-line no-await-in-loop + await new Promise(resolve => setTimeout(resolve, STREAM_TOKEN_DELAY)); + } + })(); + return true; }; @@ -1123,12 +1129,27 @@ const plugin: JupyterFrontEndPlugin = { const newCellIndex = isNewEmptyNotebook(model) ? 0 : model.cells.length - 1; + model.insertCell(newCellIndex, { cell_type: 'markdown', metadata: { trusted: true }, - source: args.source as string + source: '' }); + const sourceStr = args.source as string; + const cell = np.content.widgets[newCellIndex]; + let currentText = ''; + (async () => { + for (let i = 0; i < sourceStr.length; i += STREAM_TOKEN_LEN) { + currentText += sourceStr.slice(i, i + STREAM_TOKEN_LEN); + cell.model.sharedModel.source = currentText; + // Wait for a short delay to simulate streaming + // eslint-disable-next-line no-await-in-loop + await new Promise(resolve => + setTimeout(resolve, STREAM_TOKEN_DELAY) + ); + } + })(); return true; } }); @@ -1145,12 +1166,28 @@ const plugin: JupyterFrontEndPlugin = { const newCellIndex = isNewEmptyNotebook(model) ? 0 : model.cells.length - 1; + model.insertCell(newCellIndex, { cell_type: 'code', metadata: { trusted: true }, - source: args.source as string + source: '' }); + const sourceStr = args.source as string; + const cell = np.content.widgets[newCellIndex]; + let currentText = ''; + (async () => { + for (let i = 0; i < sourceStr.length; i += STREAM_TOKEN_LEN) { + currentText += sourceStr.slice(i, i + STREAM_TOKEN_LEN); + cell.model.sharedModel.source = currentText; + // Wait for a short delay to simulate streaming + // eslint-disable-next-line no-await-in-loop + await new Promise(resolve => + setTimeout(resolve, STREAM_TOKEN_DELAY) + ); + } + })(); + return true; } });