diff --git a/api/routers/agent.py b/api/routers/agent.py index 3eeb2a73..a034d5de 100644 --- a/api/routers/agent.py +++ b/api/routers/agent.py @@ -420,6 +420,25 @@ async def list_ollama_models(ollama_url: str = "http://localhost:11434"): return {"models": []} +async def _unload_llm_after_workflow( + client: httpx.AsyncClient, + request: AgentChatRequest, + actions_done: list[ActionDone], +) -> None: + """Free the Ollama model's VRAM once a workflow has been dispatched, so the + workflow gets the full GPU. Best-effort — never fail the chat over it.""" + if not any(a.tool == "run_workflow" for a in actions_done): + return + try: + await client.post( + f"{request.ollama_url}/api/generate", + json={"model": request.model, "keep_alive": 0}, + timeout=5.0, + ) + except Exception: + pass + + @router.post("/chat", response_model=AgentChatResponse) async def agent_chat(request: AgentChatRequest): messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}] @@ -488,6 +507,7 @@ async def agent_chat(request: AgentChatRequest): tool_calls = msg.get("tool_calls") or [] if not tool_calls: + await _unload_llm_after_workflow(client, request, actions_done) combined_thinking = "\n\n---\n\n".join(all_thinking) if all_thinking else None return AgentChatResponse( message=clean_content, @@ -501,17 +521,8 @@ async def agent_chat(request: AgentChatRequest): actions_done.append(ActionDone(tool=fn["name"], result=result_text, payload=payload)) messages.append({"role": "tool", "content": result_text}) - has_workflow = any(a.tool == "run_workflow" for a in actions_done) - if has_workflow: - # Unload LLM from VRAM immediately so the workflow has full GPU memory - try: - await client.post( - f"{request.ollama_url}/api/generate", - json={"model": request.model, "keep_alive": 0}, - timeout=5.0, - ) - except Exception: - pass + # Loop exhausted without a final answer: still free VRAM if a workflow ran. + await _unload_llm_after_workflow(client, request, actions_done) combined_thinking = "\n\n---\n\n".join(all_thinking) if all_thinking else None return AgentChatResponse(message="Reached maximum tool iterations.", actions=actions_done, thinking=combined_thinking) diff --git a/api/tests/test_agent_workflow_unload.py b/api/tests/test_agent_workflow_unload.py new file mode 100644 index 00000000..d5d7724c --- /dev/null +++ b/api/tests/test_agent_workflow_unload.py @@ -0,0 +1,77 @@ +import asyncio +import json +import unittest +from unittest import mock + +import httpx + +import routers.agent as agent + + +class _ScriptedOllama: + """MockTransport wiring: serves a scripted /api/chat sequence and records + every /api/generate (VRAM keep_alive) call agent_chat makes.""" + + def __init__(self, chat_bodies) -> None: + self._chat = list(chat_bodies) + self.generate_calls: list[dict] = [] + self._real = httpx.AsyncClient + + def _handler(self, request: httpx.Request) -> httpx.Response: + path = request.url.path + if path == "/api/chat": + return httpx.Response(200, json=self._chat.pop(0)) + if path == "/api/generate": + self.generate_calls.append(json.loads(request.content)) + return httpx.Response(200, json={}) + return httpx.Response(404, json={}) + + def __call__(self, *args, **kwargs): + kwargs["transport"] = httpx.MockTransport(self._handler) + return self._real(*args, **kwargs) + + +def _assistant(content: str = "", tool_calls=None) -> dict: + msg: dict = {"role": "assistant", "content": content} + if tool_calls: + msg["tool_calls"] = tool_calls + return {"message": msg} + + +def _run(chat_bodies, context) -> _ScriptedOllama: + scripted = _ScriptedOllama(chat_bodies) + request = agent.AgentChatRequest( + messages=[agent.ChatMessage(role="user", content="make me a thing")], + ollama_url="http://ollama.test", + model="llama-test", + context=context, + ) + with mock.patch.object(agent.httpx, "AsyncClient", scripted): + asyncio.run(agent.agent_chat(request)) + return scripted + + +class WorkflowVramUnloadTests(unittest.TestCase): + def test_llm_unloaded_on_the_normal_return_after_a_workflow(self) -> None: + # Round 1 dispatches a workflow; round 2 is the final answer (no tools) and + # returns early. The VRAM unload must still fire — before the fix it lived + # after the loop and this common path skipped it entirely. + chat = [ + _assistant(tool_calls=[ + {"function": {"name": "run_workflow", "arguments": {"workflow_id": "wf1"}}}, + ]), + _assistant(content="Running your workflow now."), + ] + scripted = _run(chat, {"workflows": [{"id": "wf1", "name": "My Workflow"}]}) + self.assertTrue( + any(call.get("keep_alive") == 0 for call in scripted.generate_calls), + f"expected a keep_alive:0 unload, got {scripted.generate_calls}", + ) + + def test_no_unload_when_no_workflow_was_dispatched(self) -> None: + scripted = _run([_assistant(content="Here is some info.")], {}) + self.assertEqual(scripted.generate_calls, []) + + +if __name__ == "__main__": + unittest.main()