Allow specifying models for some or all agents.

lynxkite-lynxscribe/src/lynxkite_lynxscribe/agentic.py

@@ -22,13 +22,19 @@ op = ops.op_registration(
 
 @op("Chat frontend", color="gray", outputs=[], view="service")
 def chat_frontend(agent: dict):
+    return agent_from_dict(agent, default_model={"name": "openai", "model_name": "gpt-4.1-nano"})
+
+
+def agent_from_dict(agent: dict, default_model: dict) -> "Agent":
     agent = agent["agent"]
+    model = agent.get("model", default_model)
     return Agent(
-        agent["name"],
-        agent["description"],
-        agent["system_prompt"],
-        agent["mcp_servers"],
-        [chat_frontend(a) for a in agent["sub_agents"]],
+        name=agent["name"],
+        description=agent["description"],
+        system_prompt=agent["system_prompt"],
+        mcp_servers=agent["mcp_servers"],
+        model=model,
+        agents=[agent_from_dict(a, default_model=model) for a in agent["sub_agents"]],
     )
 
 
@@ -40,11 +46,19 @@ def agent(
     description: ops.LongStr = "This agent helps with various tasks.",
     system_prompt: ops.LongStr = "You are a helpful assistant.",
 ):
+    """You can connect tools, other agents, or models to the input of this agent.
+
+    If a model is connected, it has the be the only model. This agent and all its
+    sub-agents will use this model.
+    """
     prompt = [system_prompt]
+    models = []
     for tool in tools:
-        if tool.get("extra_prompt"):
+        if "model" in tool:
+            models.append(tool["llm"]["model"])
+        elif tool.get("extra_prompt"):
             prompt.append(tool["extra_prompt"])
-    return {
+    params = {
         "agent": {
             "name": name,
             "description": description,
@@ -53,6 +67,19 @@ def agent(
             "sub_agents": [t for t in tools if "agent" in t],
         }
     }
+    if models:
+        assert len(models) == 1, "Only one model can be connected to an agent."
+        params["agent"]["model"] = models[0]
+    return params
+
+
+@op("Model: OpenAI", color="orange")
+def openai_model(*, model_name: str, base_url: str = ""):
+    """Use the OPENAI_API_KEY environment variable to provide authentication."""
+    params = {"model": {"name": "openai", "model_name": model_name}}
+    if base_url:
+        params["model"]["base_url"] = base_url
+    return params
 
 
 @op("MCP: Custom", color="green")
@@ -104,6 +131,7 @@ class Agent:
         prompt: str,
         mcp_servers: list[list[str]],
         agents: list["Agent"],
+        model: dict,
     ):
         self.name = name
         self.description = description
@@ -112,6 +140,8 @@ class Agent:
         self.agents = agents
         self.mcp_client = None
         self.task_solver = None
+        self.llm_engine_params = {**model}
+        self.model_name = self.llm_engine_params.pop("model_name")
 
     async def init(self):
         if self.task_solver is not None:
@@ -120,8 +150,8 @@ class Agent:
         await self.mcp_client.aenter()
         agents_as_functions = [agent.as_function() for agent in self.agents]
         self.task_solver = TaskSolver(
-            llm=get_llm_engine(name="openai"),
-            model="gpt-4.1-nano",
+            llm=get_llm_engine(**self.llm_engine_params),
+            model=self.model_name,
             initial_messages=[self.prompt],
             functions=[*self.mcp_client.functions, *agents_as_functions],
             tool_choice="required",