Skip to content

Use sys.executable for Python commands and add LLM config options #165

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
DomPTech wants to merge 9 commits into
base: main
Choose a base branch
from
Draft

Use sys.executable for Python commands and add LLM config options #165

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
2f42cbf
fix: use sys.executable for Python commands in `run_servers.py` inste…
DomPTech Aug 7, 2026
3323636
feat: new LLM config options for APIs and made ollama model config ac…
DomPTech Aug 7, 2026
e3cd8f3
Merge branch 'pycroscopy:main' into main
DomPTech Aug 8, 2026
fefaeea
Merge branch 'pycroscopy:main' into main
DomPTech Aug 8, 2026
75d5f52
Merge branch 'pycroscopy:main' into main
DomPTech Aug 10, 2026
c193a61
docs: Added LLM Device doc
DomPTech Aug 10, 2026
f6430a8
fix: replace mcp_url property in LLM Device with mcp_config to suppor…
DomPTech Aug 10, 2026
708b511
feat: LLM Device uses a single, updating instance of a MultiServerMCP…
DomPTech Aug 10, 2026
3ca2cc7
Merge branch 'pycroscopy:main' into main
DomPTech Aug 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 58 additions & 29 deletions asyncroscopy/mcp/llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,11 +51,21 @@ class AgentState(TypedDict):


class LLM(Device):
mcp_url = device_property(dtype=str, default_value="http://127.0.0.1:8000/mcp")
startup_agents = device_property(dtype=(str,), default_value=())
ollama_model = device_property(dtype=str, default_value="gemma4:31b")
use_init_chat_model = device_property(dtype=bool, default_value=False)
model_provider = device_property(dtype=str, default_value="ollama")
mcp_config = device_property(dtype=str, doc="JSON-serialized config of the MCP server to initially connect to: {'url': '...', 'transport': '...'}.")
startup_agents = device_property(dtype=(str,), default_value=(), doc="List of JSON-serialized Agent configs to spawn on startup.")

# Provider selection
use_init_chat_model = device_property(dtype=bool, default_value=False, doc="If true, use the init_chat_model function to initialize the LLM.")
model_provider = device_property(dtype=str, default_value="ollama", doc="The model provider to use for the LLM. Options: 'ollama', 'openai', etc.")

# Generic init_chat_model config
chat_model_name = device_property(dtype=str, default_value="gpt-4o", doc="The name of the chat model to use for the LLM")
api_key = device_property(dtype=str, default_value="", doc="The API key for the model provider")
api_base = device_property(dtype=str, default_value="", doc="The base URL for the API")

@utkarshp1161 utkarshp1161 Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If chat chat_model_name has default value gpt-4o the api_base should be https://api.openai.com/v1?

@DomPTech DomPTech Aug 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

init_chat_model doesn't require an api_base for most providers (for example if you wanted to use openai you just set the provider to openai and the model name to gpt-4o and it automatically handles the routing).


# Ollama config
ollama_model = device_property(dtype=str, default_value="gemma4:31b", doc="The Ollama model ID to use for the LLM")
auto_pull_model = device_property(dtype=bool, default_value=True, doc="If true, automatically pull the Ollama model if it is not already downloaded.")
Comment on lines +62 to +68

@utkarshp1161 utkarshp1161 Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder why ollama is treated here as a special case? just because it is offline?

@DomPTech DomPTech Aug 12, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct (this is the provider we tend to use the most often), so I added special functionality for auto starting ollama servers, pulling models, etc. But your comment actually did make me realize we don't need a device_property of ollama_model (since you can just do provider of ollama and chat_model_name of "gemma4:31b" example).


max_steps = attribute(label="Max Steps", dtype=int, access=tango.AttrWriteType.READ_WRITE)

Expand All @@ -69,7 +79,9 @@ async def init_device(self) -> None:
# Registries
self._agents: list[Agent] = []
self._tools: list[BaseTool] = []
self._mcp_clients: list[MultiServerMCPClient] = []

# Manages MCP connections
self._mcp_client = MultiServerMCPClient({})

if self.startup_agents:
self._agents = [Agent(**json.loads(agent_json)) for agent_json in self.startup_agents]
Expand All @@ -81,17 +93,23 @@ async def init_device(self) -> None:

if self.use_init_chat_model: # Initialize from most model providers (e.g., OpenAI)
self.info_stream("Initializing via init_chat_model")
self._model = init_chat_model(
model=self.ollama_model,
model_provider=self.model_provider,
temperature=0
)

model_kwargs = {
"model": self.chat_model_name,
"model_provider": self.model_provider
}
if self.api_key:
model_kwargs["api_key"] = self.api_key
if self.api_base:
model_kwargs["api_base"] = self.api_base

self._model = init_chat_model(**model_kwargs)
else: # Initialize locally via Ollama
from langchain_ollama import ChatOllama
self.info_stream("Initializing via ChatOllama")
self._model = ChatOllama(
model=self.ollama_model,
temperature=0,
temperature=0.7,
reasoning=False,
)

Expand All @@ -102,10 +120,9 @@ async def init_device(self) -> None:
print(f"[SYSTEM]: Model pre-warmed in {time.time() - start_warmup:.2f}s!")

# Connect to an MCP server initially if specified
if self.mcp_url:
config = json.dumps({"url": self.mcp_url, "transport": "streamable_http"})
if not await self.ConnectMCP(config):
print(f"[SYSTEM]: Failed to connect to MCP Server at {self.mcp_url}.")
if self.mcp_config:
if not await self.ConnectMCP(self.mcp_config):
print(f"[SYSTEM]: Failed to connect to MCP Server: {self.mcp_config}.")

self.set_state(tango.DevState.ON)
except Exception as e:
Expand All @@ -126,8 +143,12 @@ def agents(self) -> list[str]:
"""Return a list of the names of all currently spawned agents."""
return [agent.name for agent in self._agents]

@attribute(dtype=str)
def mcp_connections(self) -> str:
return str(self._mcp_client.connections)

async def ensure_ollama_running(self, host: str = "http://localhost:11434", timeout: int = 10) -> None:
"""Check if Ollama server is running, offloaded to prevent blocking the Tango loop."""
"""Check if Ollama server is running, starting it and downloading the model if necessary."""

def _sync_check():
tags_url = f"{host.rstrip('/')}/api/tags"
Expand All @@ -136,8 +157,19 @@ def _sync_check():
return
except (urllib.error.URLError, TimeoutError, ConnectionRefusedError):
pass

try:

if self.auto_pull_model: # Run command to download the model if it is not already

@utkarshp1161 utkarshp1161 Aug 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would give a sort of warning message to the user with list of models already there to choose from. More like: print("this model_x is not available, these models are available .....model_a, model_b..., either choose from these or Proceed as is to pull the model_x which may take some time")

@DomPTech DomPTech Aug 12, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good thought!

print(f"[SYSTEM]: Ensuring model '{self.ollama_model}' is pulled (this may take a while if it is not already downloaded)...")
try:
subprocess.run(
["ollama", "pull", self.ollama_model],
check=True,
stderr=subprocess.DEVNULL
)
except subprocess.CalledProcessError as e:
raise RuntimeError(f"Failed to pull Ollama model {self.ollama_model}: {e}")

try: # Run command to serve the model
subprocess.Popen(
["ollama", "serve"],
stdout=subprocess.DEVNULL,
Expand Down Expand Up @@ -182,19 +214,16 @@ async def ConnectMCP(self, config: str) -> bool:
"""Connect to an MCP server and inherit its tools. Returns true for success."""
try:
args = json.loads(config)
url = args.get("url")
transport = args.get("transport", "streamable_http")

server_id = f"server_{len(self._mcp_clients)}"
client = MultiServerMCPClient({server_id: {"url": url, "transport": transport}})
# Generate unique ID for this connection and add it to the client
server_id = f"server_{len(self._mcp_client.connections)}"
self._mcp_client.connections[server_id] = args

print(f"\n[SYSTEM]: Connecting to MCP Server at {url}...")
print(f"\n[SYSTEM]: Added MCP server {server_id}. Fetching tools...")

tools = await client.get_tools()
self._tools = await self._mcp_client.get_tools()

self._mcp_clients.append(client)
self._tools.extend(tools)
print(f"[SYSTEM]: Connected. Inherited {len(tools)} tools.")
print(f"[SYSTEM]: Connected. Inherited {len(self._tools)} tools.")
except Exception as e:
self.error_stream(f"Failed to connect to MCP server: {e}")
return False
Expand Down Expand Up @@ -245,7 +274,7 @@ def _extract_json(self, text: str) -> str:
return text.strip()

def _parse_routing_decision(self, content: str, valid_options: list[str], fallback: str) -> tuple[str, str]:
"""Parse a supervisor response's {'next': ...} decision, falling back on any error or invalid value."""
"""Parse a supervisor response's {'next': ...} decision with a fallback."""
try:
decision = json.loads(self._extract_json(content))
next_agent = decision.get("next", fallback)
Expand Down
6 changes: 4 additions & 2 deletions configs/gemma-llm.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,10 @@ tango:
host: localhost
port: 9094

mcp_url: "http://127.0.0.1:8001/mcp"
local_model_path: C:\Users\Public\Desktop\Agents\Gemma4-31B-4bit
mcp_config:
url: "http://127.0.0.1:8000/mcp"
transport: streamable-http
ollama_model: gemma4:31b

startup_agents:
- name: "base"
Expand Down
2 changes: 1 addition & 1 deletion configs/mcp_dt.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
# machine/IP. Set mcp.http_host to 0.0.0.0 when other machines need to connect.

tango:
host: 127.0.0.1
host: localhost
port: 9094

mcp:
Expand Down
36 changes: 36 additions & 0 deletions docs/MCP/llm_device.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
# LLM Device

The `LLM` device is a Tango device that wraps an LangChain AI agent swarm. It enables AI agents to interact with the hardware tools exposed via the MCP server. It is started with `startup_scripts/run_llm.py`.

## Overview

The device initializes a LangChain/LangGraph-based swarm. It exposes Tango Commands to connect to multiple MCP servers, spawn specialized worker agents, and take queries. A central Supervisor agent routes tasks to the appropriate worker based on user input.

This Device may be initialized with a local model, via Ollama, or more generally via Langchain's `init_chat_model` method that takes a model and provider, and optionally other args/kwargs like an API key. With Ollama, the model is automatically served if the server isn't already running and optionally pulled if not already downloaed.

## Commands

- **`SpawnAgent(config: str)`**
Creates a new worker agent. Expects a JSON-serialized string:
```json
{
"name": "agent_name",
"system_prompt": "Agent role and instructions.",
"description": "An optional description to help the Supervisor with routing",
"model": "optional_model_override",
"tools": ["glob_pattern_1", "glob_pattern_2"]
}
```
Glob patterns can be used for giving agents multiple related tools, such as "*image" to give it all tools whose names end in the word "image".

- **`ConnectMCP(config: str)`**
Connects to an MCP server and inherits its tools. Expects a JSON string:
```json
{
"url": "http://127.0.0.1:8000/mcp",
"transport": "streamable_http"
}
```

- **`Query(prompt: str)`**
Queries the swarm. If more than one agent is in the swarm, a Supervisor routes this prompt to the appropriate agent. Otherwise, the single agent handles the request like a normal chatbot.
1 change: 1 addition & 0 deletions docs/index.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ Use this site to navigate contributor guidance, microscope architecture notes, h
- [Add a Detector](Adding_New_Hardware/add_detector.md): detector onboarding checklist and implementation notes.
- [Data Integration (Tiled)](Tiled_server/data_integration.md): how acquisitions are saved, registered, and served via the DATA device and Tiled.
- [MCP Server Documentation](MCP/mcp_server.md): how Tango commands are exposed to MCP-compatible agents.
- [LLM Device Documentation](MCP/llm_device.md): configuring the agent swarm and MCP tool integration.

## Operation

Expand Down
43 changes: 13 additions & 30 deletions notebooks/11_Test_AI_Agent.ipynb
View file
Open in desktop

@utkarshp1161 utkarshp1161 Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just reading the notebook :
I wonder if below means the agent was not able to call the right tools and then failed? Will increasing Recursion limit help?

prompt = "Get a scanned haadf image. Then get an EDS spectrum."
response = llm.query(prompt)
print(response)

Recursion limit of 5 reached without hitting a stop condition. You can increase the limit by setting the recursion_limit config key.
For troubleshooting, visit: https://docs.langchain.com/oss/python/langgraph/errors/GRAPH_RECURSION_LIMIT

@DomPTech DomPTech Aug 12, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was a very interesting thing during testing actually. Basically what happened was the agent successfully acquired the image, and then tried to acquire an EDS spectrum. However, for some reason the acquire_spectrum tool on the DigitalTwin always seems to fail with a Tango error (according to the stack trace because the EDS device is not properly initialized). Thus it continued to try different things to no avail, even though it technically was correct. Perhaps we should raise an issue for this.

@utkarshp1161 utkarshp1161 Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting. A issue on this would be nice.

Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 18,
"id": "c4d20080",
"metadata": {},
"outputs": [],
Expand All @@ -65,7 +65,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 19,
"id": "04f5aae1",
"metadata": {},
"outputs": [
Expand All @@ -76,22 +76,9 @@
"localhost:9094\n",
"asyncroscopy/instrument/default ON\n",
"asyncroscopy/data/default ON\n",
"asyncroscopy/eds/default ON\n",
"asyncroscopy/llm/default ON\n",
"Tiled: http://localhost:9091\n"
]
},
{
"ename": "AttributeError",
"evalue": "acquire_spectrum",
"output_type": "error",
"traceback": [
"\u001b[31m---------------------------------------------------------------------------\u001b[39m",
"\u001b[31mAttributeError\u001b[39m Traceback (most recent call last)",
"\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[22]\u001b[39m\u001b[32m, line 21\u001b[39m\n\u001b[32m 17\u001b[39m tiled_config = json.loads(data.get_config())\n\u001b[32m 18\u001b[39m client = from_uri(tiled_config[\u001b[33m\"uri\"\u001b[39m])\n\u001b[32m 19\u001b[39m print(\u001b[33m\"Tiled:\"\u001b[39m, tiled_config[\u001b[33m\"uri\"\u001b[39m])\n\u001b[32m 20\u001b[39m \n\u001b[32m---> \u001b[39m\u001b[32m21\u001b[39m eds.acquire_spectrum(\u001b[33m\"EDS\"\u001b[39m)\n",
"\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\dpelaia\\Documents\\asyncroscopy-agent-version\\.venv\\Lib\\site-packages\\tango\\device_proxy.py:2319\u001b[39m, in \u001b[36m__safe_call.<locals>.safe_call_wrapper\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m 2317\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m._initialized:\n\u001b[32m 2318\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[33m\"\u001b[39m\u001b[33mDeviceProxy object was not fully initialized.\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m-> \u001b[39m\u001b[32m2319\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mfn\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43mself\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n",
"\u001b[36mFile \u001b[39m\u001b[32mc:\\Users\\dpelaia\\Documents\\asyncroscopy-agent-version\\.venv\\Lib\\site-packages\\tango\\device_proxy.py:484\u001b[39m, in \u001b[36m__DeviceProxy__getattr\u001b[39m\u001b[34m(self, name)\u001b[39m\n\u001b[32m 481\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m attr_info:\n\u001b[32m 482\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m __get_attribute_value(\u001b[38;5;28mself\u001b[39m, attr_info, name)\n\u001b[32m--> \u001b[39m\u001b[32m484\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m(name) \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mcause\u001b[39;00m\n\u001b[32m 485\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m 486\u001b[39m \u001b[38;5;28;01mdel\u001b[39;00m cause\n",
"\u001b[31mAttributeError\u001b[39m: acquire_spectrum"
]
}
],
"source": [
Expand All @@ -117,7 +104,7 @@
},
{
"cell_type": "code",
"execution_count": 17,
"execution_count": 20,
"id": "ce4eba9c",
"metadata": {},
"outputs": [
Expand All @@ -127,7 +114,7 @@
"True"
]
},
"execution_count": 17,
"execution_count": 20,
"metadata": {},
"output_type": "execute_result"
}
Expand All @@ -147,20 +134,22 @@
},
{
"cell_type": "code",
"execution_count": 19,
"execution_count": 21,
"id": "1e1e508a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"('base', 'image', 'eds')\n"
"('base', 'image', 'eds')\n",
"{'server_0': {'url': 'http://127.0.0.1:8000/mcp', 'transport': 'streamable-http'}}\n"
]
}
],
"source": [
"print(llm.agents)"
"print(llm.agents)\n",
"print(llm.mcp_connections)"
]
},
{
Expand All @@ -173,22 +162,16 @@
},
{
"cell_type": "code",
"execution_count": 20,
"execution_count": 22,
"id": "a1e65cdd",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"The available devices are:\n",
"- asyncroscopy/camera/default\n",
"- asyncroscopy/data/default\n",
"- asyncroscopy/eds/default\n",
"- asyncroscopy/instrument/default\n",
"- asyncroscopy/llm/default\n",
"- asyncroscopy/scan/default\n",
"- asyncroscopy/stage/default\n"
"Recursion limit of 5 reached without hitting a stop condition. You can increase the limit by setting the `recursion_limit` config key.\n",
"For troubleshooting, visit: https://docs.langchain.com/oss/python/langgraph/errors/GRAPH_RECURSION_LIMIT\n"
]
}
],
Expand Down
32 changes: 17 additions & 15 deletions startup_scripts/run_llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from typing import Literal
import json
import sys
import argparse
Expand Down Expand Up @@ -27,25 +28,24 @@ class TangoConfig:
host: str
port: int

@dataclass
class MCPConfig:
url: str | None = None
transport: Literal["stdio", "http", "sse", "streamable-http"] = "streamable-http"

@dataclass
class LLMConfig:
tango: TangoConfig
mcp_url: str
local_model_path: str | None = None
model_provider: str | None = None
model_name: str | None = None
mcp_config: MCPConfig

chat_model_name: str | None = None
api_key: str | None = None
startup_agents: list[Agent] | None = None
api_base: str | None = None

def __post_init__(self):
# Convert tango dict to TangoConfig
if isinstance(self.tango, dict):
self.tango = TangoConfig(**self.tango)
ollama_model: str | None = None
auto_pull_model: bool = True

def _require(mapping: dict, key: str, where: str):
if not isinstance(mapping, dict) or key not in mapping:
raise KeyError(f"Config section '{where}' is missing required key '{key}'")
return mapping[key]
startup_agents: list[Agent] | None = None


def load_config(path: Path) -> LLMConfig:
Expand Down Expand Up @@ -73,6 +73,8 @@ def register_device(config: LLMConfig | None):
properties[key] = value
if key == "startup_agents":
properties[key] = [json.dumps(agent) for agent in value]
elif key == "mcp_config":
properties[key] = json.dumps(value)

database.put_device_property(DEVICE_NAME, properties)
print(f"Set device properties: {properties}")
Expand All @@ -91,12 +93,12 @@ def main(argv: list[str] | None = None) -> int:
print(f'Config error: {exc}', file=sys.stderr)
return 1

tango_host = f'{config.tango.host}:{config.tango.port}'
tango_host = f'{config.tango["host"]}:{config.tango["port"]}'
os.environ['TANGO_HOST'] = tango_host

register_device(config)

command = ["uv", "run", "python", "-m", "asyncroscopy.mcp.llm", INSTANCE_NAME]
command = [sys.executable, "-m", "asyncroscopy.mcp.llm", INSTANCE_NAME]
env = {**os.environ, 'TANGO_HOST': tango_host, 'PYTHONUNBUFFERED': '1'}

try:
Expand Down
Loading
Loading