agno 2.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (314) hide show
  1. agno/agent/agent.py +5540 -2273
  2. agno/api/api.py +2 -0
  3. agno/api/os.py +1 -1
  4. agno/compression/__init__.py +3 -0
  5. agno/compression/manager.py +247 -0
  6. agno/culture/__init__.py +3 -0
  7. agno/culture/manager.py +956 -0
  8. agno/db/async_postgres/__init__.py +3 -0
  9. agno/db/base.py +689 -6
  10. agno/db/dynamo/dynamo.py +933 -37
  11. agno/db/dynamo/schemas.py +174 -10
  12. agno/db/dynamo/utils.py +63 -4
  13. agno/db/firestore/firestore.py +831 -9
  14. agno/db/firestore/schemas.py +51 -0
  15. agno/db/firestore/utils.py +102 -4
  16. agno/db/gcs_json/gcs_json_db.py +660 -12
  17. agno/db/gcs_json/utils.py +60 -26
  18. agno/db/in_memory/in_memory_db.py +287 -14
  19. agno/db/in_memory/utils.py +60 -2
  20. agno/db/json/json_db.py +590 -14
  21. agno/db/json/utils.py +60 -26
  22. agno/db/migrations/manager.py +199 -0
  23. agno/db/migrations/v1_to_v2.py +43 -13
  24. agno/db/migrations/versions/__init__.py +0 -0
  25. agno/db/migrations/versions/v2_3_0.py +938 -0
  26. agno/db/mongo/__init__.py +15 -1
  27. agno/db/mongo/async_mongo.py +2760 -0
  28. agno/db/mongo/mongo.py +879 -11
  29. agno/db/mongo/schemas.py +42 -0
  30. agno/db/mongo/utils.py +80 -8
  31. agno/db/mysql/__init__.py +2 -1
  32. agno/db/mysql/async_mysql.py +2912 -0
  33. agno/db/mysql/mysql.py +946 -68
  34. agno/db/mysql/schemas.py +72 -10
  35. agno/db/mysql/utils.py +198 -7
  36. agno/db/postgres/__init__.py +2 -1
  37. agno/db/postgres/async_postgres.py +2579 -0
  38. agno/db/postgres/postgres.py +942 -57
  39. agno/db/postgres/schemas.py +81 -18
  40. agno/db/postgres/utils.py +164 -2
  41. agno/db/redis/redis.py +671 -7
  42. agno/db/redis/schemas.py +50 -0
  43. agno/db/redis/utils.py +65 -7
  44. agno/db/schemas/__init__.py +2 -1
  45. agno/db/schemas/culture.py +120 -0
  46. agno/db/schemas/evals.py +1 -0
  47. agno/db/schemas/memory.py +17 -2
  48. agno/db/singlestore/schemas.py +63 -0
  49. agno/db/singlestore/singlestore.py +949 -83
  50. agno/db/singlestore/utils.py +60 -2
  51. agno/db/sqlite/__init__.py +2 -1
  52. agno/db/sqlite/async_sqlite.py +2911 -0
  53. agno/db/sqlite/schemas.py +62 -0
  54. agno/db/sqlite/sqlite.py +965 -46
  55. agno/db/sqlite/utils.py +169 -8
  56. agno/db/surrealdb/__init__.py +3 -0
  57. agno/db/surrealdb/metrics.py +292 -0
  58. agno/db/surrealdb/models.py +334 -0
  59. agno/db/surrealdb/queries.py +71 -0
  60. agno/db/surrealdb/surrealdb.py +1908 -0
  61. agno/db/surrealdb/utils.py +147 -0
  62. agno/db/utils.py +2 -0
  63. agno/eval/__init__.py +10 -0
  64. agno/eval/accuracy.py +75 -55
  65. agno/eval/agent_as_judge.py +861 -0
  66. agno/eval/base.py +29 -0
  67. agno/eval/performance.py +16 -7
  68. agno/eval/reliability.py +28 -16
  69. agno/eval/utils.py +35 -17
  70. agno/exceptions.py +27 -2
  71. agno/filters.py +354 -0
  72. agno/guardrails/prompt_injection.py +1 -0
  73. agno/hooks/__init__.py +3 -0
  74. agno/hooks/decorator.py +164 -0
  75. agno/integrations/discord/client.py +1 -1
  76. agno/knowledge/chunking/agentic.py +13 -10
  77. agno/knowledge/chunking/fixed.py +4 -1
  78. agno/knowledge/chunking/semantic.py +9 -4
  79. agno/knowledge/chunking/strategy.py +59 -15
  80. agno/knowledge/embedder/fastembed.py +1 -1
  81. agno/knowledge/embedder/nebius.py +1 -1
  82. agno/knowledge/embedder/ollama.py +8 -0
  83. agno/knowledge/embedder/openai.py +8 -8
  84. agno/knowledge/embedder/sentence_transformer.py +6 -2
  85. agno/knowledge/embedder/vllm.py +262 -0
  86. agno/knowledge/knowledge.py +1618 -318
  87. agno/knowledge/reader/base.py +6 -2
  88. agno/knowledge/reader/csv_reader.py +8 -10
  89. agno/knowledge/reader/docx_reader.py +5 -6
  90. agno/knowledge/reader/field_labeled_csv_reader.py +16 -20
  91. agno/knowledge/reader/json_reader.py +5 -4
  92. agno/knowledge/reader/markdown_reader.py +8 -8
  93. agno/knowledge/reader/pdf_reader.py +17 -19
  94. agno/knowledge/reader/pptx_reader.py +101 -0
  95. agno/knowledge/reader/reader_factory.py +32 -3
  96. agno/knowledge/reader/s3_reader.py +3 -3
  97. agno/knowledge/reader/tavily_reader.py +193 -0
  98. agno/knowledge/reader/text_reader.py +22 -10
  99. agno/knowledge/reader/web_search_reader.py +1 -48
  100. agno/knowledge/reader/website_reader.py +10 -10
  101. agno/knowledge/reader/wikipedia_reader.py +33 -1
  102. agno/knowledge/types.py +1 -0
  103. agno/knowledge/utils.py +72 -7
  104. agno/media.py +22 -6
  105. agno/memory/__init__.py +14 -1
  106. agno/memory/manager.py +544 -83
  107. agno/memory/strategies/__init__.py +15 -0
  108. agno/memory/strategies/base.py +66 -0
  109. agno/memory/strategies/summarize.py +196 -0
  110. agno/memory/strategies/types.py +37 -0
  111. agno/models/aimlapi/aimlapi.py +17 -0
  112. agno/models/anthropic/claude.py +515 -40
  113. agno/models/aws/bedrock.py +102 -21
  114. agno/models/aws/claude.py +131 -274
  115. agno/models/azure/ai_foundry.py +41 -19
  116. agno/models/azure/openai_chat.py +39 -8
  117. agno/models/base.py +1249 -525
  118. agno/models/cerebras/cerebras.py +91 -21
  119. agno/models/cerebras/cerebras_openai.py +21 -2
  120. agno/models/cohere/chat.py +40 -6
  121. agno/models/cometapi/cometapi.py +18 -1
  122. agno/models/dashscope/dashscope.py +2 -3
  123. agno/models/deepinfra/deepinfra.py +18 -1
  124. agno/models/deepseek/deepseek.py +69 -3
  125. agno/models/fireworks/fireworks.py +18 -1
  126. agno/models/google/gemini.py +877 -80
  127. agno/models/google/utils.py +22 -0
  128. agno/models/groq/groq.py +51 -18
  129. agno/models/huggingface/huggingface.py +17 -6
  130. agno/models/ibm/watsonx.py +16 -6
  131. agno/models/internlm/internlm.py +18 -1
  132. agno/models/langdb/langdb.py +13 -1
  133. agno/models/litellm/chat.py +44 -9
  134. agno/models/litellm/litellm_openai.py +18 -1
  135. agno/models/message.py +28 -5
  136. agno/models/meta/llama.py +47 -14
  137. agno/models/meta/llama_openai.py +22 -17
  138. agno/models/mistral/mistral.py +8 -4
  139. agno/models/nebius/nebius.py +6 -7
  140. agno/models/nvidia/nvidia.py +20 -3
  141. agno/models/ollama/chat.py +24 -8
  142. agno/models/openai/chat.py +104 -29
  143. agno/models/openai/responses.py +101 -81
  144. agno/models/openrouter/openrouter.py +60 -3
  145. agno/models/perplexity/perplexity.py +17 -1
  146. agno/models/portkey/portkey.py +7 -6
  147. agno/models/requesty/requesty.py +24 -4
  148. agno/models/response.py +73 -2
  149. agno/models/sambanova/sambanova.py +20 -3
  150. agno/models/siliconflow/siliconflow.py +19 -2
  151. agno/models/together/together.py +20 -3
  152. agno/models/utils.py +254 -8
  153. agno/models/vercel/v0.py +20 -3
  154. agno/models/vertexai/__init__.py +0 -0
  155. agno/models/vertexai/claude.py +190 -0
  156. agno/models/vllm/vllm.py +19 -14
  157. agno/models/xai/xai.py +19 -2
  158. agno/os/app.py +549 -152
  159. agno/os/auth.py +190 -3
  160. agno/os/config.py +23 -0
  161. agno/os/interfaces/a2a/router.py +8 -11
  162. agno/os/interfaces/a2a/utils.py +1 -1
  163. agno/os/interfaces/agui/router.py +18 -3
  164. agno/os/interfaces/agui/utils.py +152 -39
  165. agno/os/interfaces/slack/router.py +55 -37
  166. agno/os/interfaces/slack/slack.py +9 -1
  167. agno/os/interfaces/whatsapp/router.py +0 -1
  168. agno/os/interfaces/whatsapp/security.py +3 -1
  169. agno/os/mcp.py +110 -52
  170. agno/os/middleware/__init__.py +2 -0
  171. agno/os/middleware/jwt.py +676 -112
  172. agno/os/router.py +40 -1478
  173. agno/os/routers/agents/__init__.py +3 -0
  174. agno/os/routers/agents/router.py +599 -0
  175. agno/os/routers/agents/schema.py +261 -0
  176. agno/os/routers/evals/evals.py +96 -39
  177. agno/os/routers/evals/schemas.py +65 -33
  178. agno/os/routers/evals/utils.py +80 -10
  179. agno/os/routers/health.py +10 -4
  180. agno/os/routers/knowledge/knowledge.py +196 -38
  181. agno/os/routers/knowledge/schemas.py +82 -22
  182. agno/os/routers/memory/memory.py +279 -52
  183. agno/os/routers/memory/schemas.py +46 -17
  184. agno/os/routers/metrics/metrics.py +20 -8
  185. agno/os/routers/metrics/schemas.py +16 -16
  186. agno/os/routers/session/session.py +462 -34
  187. agno/os/routers/teams/__init__.py +3 -0
  188. agno/os/routers/teams/router.py +512 -0
  189. agno/os/routers/teams/schema.py +257 -0
  190. agno/os/routers/traces/__init__.py +3 -0
  191. agno/os/routers/traces/schemas.py +414 -0
  192. agno/os/routers/traces/traces.py +499 -0
  193. agno/os/routers/workflows/__init__.py +3 -0
  194. agno/os/routers/workflows/router.py +624 -0
  195. agno/os/routers/workflows/schema.py +75 -0
  196. agno/os/schema.py +256 -693
  197. agno/os/scopes.py +469 -0
  198. agno/os/utils.py +514 -36
  199. agno/reasoning/anthropic.py +80 -0
  200. agno/reasoning/gemini.py +73 -0
  201. agno/reasoning/openai.py +5 -0
  202. agno/reasoning/vertexai.py +76 -0
  203. agno/run/__init__.py +6 -0
  204. agno/run/agent.py +155 -32
  205. agno/run/base.py +55 -3
  206. agno/run/requirement.py +181 -0
  207. agno/run/team.py +125 -38
  208. agno/run/workflow.py +72 -18
  209. agno/session/agent.py +102 -89
  210. agno/session/summary.py +56 -15
  211. agno/session/team.py +164 -90
  212. agno/session/workflow.py +405 -40
  213. agno/table.py +10 -0
  214. agno/team/team.py +3974 -1903
  215. agno/tools/dalle.py +2 -4
  216. agno/tools/eleven_labs.py +23 -25
  217. agno/tools/exa.py +21 -16
  218. agno/tools/file.py +153 -23
  219. agno/tools/file_generation.py +16 -10
  220. agno/tools/firecrawl.py +15 -7
  221. agno/tools/function.py +193 -38
  222. agno/tools/gmail.py +238 -14
  223. agno/tools/google_drive.py +271 -0
  224. agno/tools/googlecalendar.py +36 -8
  225. agno/tools/googlesheets.py +20 -5
  226. agno/tools/jira.py +20 -0
  227. agno/tools/mcp/__init__.py +10 -0
  228. agno/tools/mcp/mcp.py +331 -0
  229. agno/tools/mcp/multi_mcp.py +347 -0
  230. agno/tools/mcp/params.py +24 -0
  231. agno/tools/mcp_toolbox.py +3 -3
  232. agno/tools/models/nebius.py +5 -5
  233. agno/tools/models_labs.py +20 -10
  234. agno/tools/nano_banana.py +151 -0
  235. agno/tools/notion.py +204 -0
  236. agno/tools/parallel.py +314 -0
  237. agno/tools/postgres.py +76 -36
  238. agno/tools/redshift.py +406 -0
  239. agno/tools/scrapegraph.py +1 -1
  240. agno/tools/shopify.py +1519 -0
  241. agno/tools/slack.py +18 -3
  242. agno/tools/spotify.py +919 -0
  243. agno/tools/tavily.py +146 -0
  244. agno/tools/toolkit.py +25 -0
  245. agno/tools/workflow.py +8 -1
  246. agno/tools/yfinance.py +12 -11
  247. agno/tracing/__init__.py +12 -0
  248. agno/tracing/exporter.py +157 -0
  249. agno/tracing/schemas.py +276 -0
  250. agno/tracing/setup.py +111 -0
  251. agno/utils/agent.py +938 -0
  252. agno/utils/cryptography.py +22 -0
  253. agno/utils/dttm.py +33 -0
  254. agno/utils/events.py +151 -3
  255. agno/utils/gemini.py +15 -5
  256. agno/utils/hooks.py +118 -4
  257. agno/utils/http.py +113 -2
  258. agno/utils/knowledge.py +12 -5
  259. agno/utils/log.py +1 -0
  260. agno/utils/mcp.py +92 -2
  261. agno/utils/media.py +187 -1
  262. agno/utils/merge_dict.py +3 -3
  263. agno/utils/message.py +60 -0
  264. agno/utils/models/ai_foundry.py +9 -2
  265. agno/utils/models/claude.py +49 -14
  266. agno/utils/models/cohere.py +9 -2
  267. agno/utils/models/llama.py +9 -2
  268. agno/utils/models/mistral.py +4 -2
  269. agno/utils/print_response/agent.py +109 -16
  270. agno/utils/print_response/team.py +223 -30
  271. agno/utils/print_response/workflow.py +251 -34
  272. agno/utils/streamlit.py +1 -1
  273. agno/utils/team.py +98 -9
  274. agno/utils/tokens.py +657 -0
  275. agno/vectordb/base.py +39 -7
  276. agno/vectordb/cassandra/cassandra.py +21 -5
  277. agno/vectordb/chroma/chromadb.py +43 -12
  278. agno/vectordb/clickhouse/clickhousedb.py +21 -5
  279. agno/vectordb/couchbase/couchbase.py +29 -5
  280. agno/vectordb/lancedb/lance_db.py +92 -181
  281. agno/vectordb/langchaindb/langchaindb.py +24 -4
  282. agno/vectordb/lightrag/lightrag.py +17 -3
  283. agno/vectordb/llamaindex/llamaindexdb.py +25 -5
  284. agno/vectordb/milvus/milvus.py +50 -37
  285. agno/vectordb/mongodb/__init__.py +7 -1
  286. agno/vectordb/mongodb/mongodb.py +36 -30
  287. agno/vectordb/pgvector/pgvector.py +201 -77
  288. agno/vectordb/pineconedb/pineconedb.py +41 -23
  289. agno/vectordb/qdrant/qdrant.py +67 -54
  290. agno/vectordb/redis/__init__.py +9 -0
  291. agno/vectordb/redis/redisdb.py +682 -0
  292. agno/vectordb/singlestore/singlestore.py +50 -29
  293. agno/vectordb/surrealdb/surrealdb.py +31 -41
  294. agno/vectordb/upstashdb/upstashdb.py +34 -6
  295. agno/vectordb/weaviate/weaviate.py +53 -14
  296. agno/workflow/__init__.py +2 -0
  297. agno/workflow/agent.py +299 -0
  298. agno/workflow/condition.py +120 -18
  299. agno/workflow/loop.py +77 -10
  300. agno/workflow/parallel.py +231 -143
  301. agno/workflow/router.py +118 -17
  302. agno/workflow/step.py +609 -170
  303. agno/workflow/steps.py +73 -6
  304. agno/workflow/types.py +96 -21
  305. agno/workflow/workflow.py +2039 -262
  306. {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/METADATA +201 -66
  307. agno-2.3.13.dist-info/RECORD +613 -0
  308. agno/tools/googlesearch.py +0 -98
  309. agno/tools/mcp.py +0 -679
  310. agno/tools/memori.py +0 -339
  311. agno-2.1.2.dist-info/RECORD +0 -543
  312. {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/WHEEL +0 -0
  313. {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/licenses/LICENSE +0 -0
  314. {agno-2.1.2.dist-info → agno-2.3.13.dist-info}/top_level.txt +0 -0
agno/os/router.py CHANGED
@@ -1,28 +1,18 @@
1
- import json
2
- from itertools import chain
3
- from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Dict, List, Optional, Union, cast
4
- from uuid import uuid4
1
+ from typing import TYPE_CHECKING, List, Optional, Union, cast
5
2
 
6
3
  from fastapi import (
7
4
  APIRouter,
8
5
  Depends,
9
- File,
10
- Form,
11
6
  HTTPException,
12
- Request,
13
- UploadFile,
14
- WebSocket,
15
7
  )
16
- from fastapi.responses import JSONResponse, StreamingResponse
17
- from pydantic import BaseModel
8
+ from fastapi.responses import JSONResponse
9
+ from packaging import version
18
10
 
19
11
  from agno.agent.agent import Agent
20
- from agno.exceptions import InputCheckError, OutputCheckError
21
- from agno.media import Audio, Image, Video
22
- from agno.media import File as FileMedia
23
- from agno.os.auth import get_authentication_dependency, validate_websocket_token
12
+ from agno.db.base import AsyncBaseDb
13
+ from agno.db.migrations.manager import MigrationManager
14
+ from agno.os.auth import get_authentication_dependency
24
15
  from agno.os.schema import (
25
- AgentResponse,
26
16
  AgentSummaryResponse,
27
17
  BadRequestResponse,
28
18
  ConfigResponse,
@@ -30,508 +20,21 @@ from agno.os.schema import (
30
20
  InternalServerErrorResponse,
31
21
  Model,
32
22
  NotFoundResponse,
33
- TeamResponse,
34
23
  TeamSummaryResponse,
35
24
  UnauthenticatedResponse,
36
25
  ValidationErrorResponse,
37
- WorkflowResponse,
38
26
  WorkflowSummaryResponse,
39
27
  )
40
28
  from agno.os.settings import AgnoAPISettings
41
29
  from agno.os.utils import (
42
- get_agent_by_id,
43
- get_team_by_id,
44
- get_workflow_by_id,
45
- process_audio,
46
- process_document,
47
- process_image,
48
- process_video,
30
+ get_db,
49
31
  )
50
- from agno.run.agent import RunErrorEvent, RunOutput, RunOutputEvent
51
- from agno.run.team import RunErrorEvent as TeamRunErrorEvent
52
- from agno.run.team import TeamRunOutputEvent
53
- from agno.run.workflow import WorkflowErrorEvent, WorkflowRunOutput, WorkflowRunOutputEvent
54
32
  from agno.team.team import Team
55
- from agno.utils.log import log_debug, log_error, log_warning, logger
56
- from agno.workflow.workflow import Workflow
57
33
 
58
34
  if TYPE_CHECKING:
59
35
  from agno.os.app import AgentOS
60
36
 
61
37
 
62
- async def _get_request_kwargs(request: Request, endpoint_func: Callable) -> Dict[str, Any]:
63
- """Given a Request and an endpoint function, return a dictionary with all extra form data fields.
64
- Args:
65
- request: The FastAPI Request object
66
- endpoint_func: The function exposing the endpoint that received the request
67
-
68
- Returns:
69
- A dictionary of kwargs
70
- """
71
- import inspect
72
-
73
- form_data = await request.form()
74
- sig = inspect.signature(endpoint_func)
75
- known_fields = set(sig.parameters.keys())
76
- kwargs = {key: value for key, value in form_data.items() if key not in known_fields}
77
-
78
- # Handle JSON parameters. They are passed as strings and need to be deserialized.
79
- if session_state := kwargs.get("session_state"):
80
- try:
81
- session_state_dict = json.loads(session_state) # type: ignore
82
- kwargs["session_state"] = session_state_dict
83
- except json.JSONDecodeError:
84
- kwargs.pop("session_state")
85
- log_warning(f"Invalid session_state parameter couldn't be loaded: {session_state}")
86
-
87
- if dependencies := kwargs.get("dependencies"):
88
- try:
89
- dependencies_dict = json.loads(dependencies) # type: ignore
90
- kwargs["dependencies"] = dependencies_dict
91
- except json.JSONDecodeError:
92
- kwargs.pop("dependencies")
93
- log_warning(f"Invalid dependencies parameter couldn't be loaded: {dependencies}")
94
-
95
- return kwargs
96
-
97
-
98
- def format_sse_event(event: Union[RunOutputEvent, TeamRunOutputEvent, WorkflowRunOutputEvent]) -> str:
99
- """Parse JSON data into SSE-compliant format.
100
-
101
- Args:
102
- event_dict: Dictionary containing the event data
103
-
104
- Returns:
105
- SSE-formatted response:
106
-
107
- ```
108
- event: EventName
109
- data: { ... }
110
-
111
- event: AnotherEventName
112
- data: { ... }
113
- ```
114
- """
115
- try:
116
- # Parse the JSON to extract the event type
117
- event_type = event.event or "message"
118
-
119
- # Serialize to valid JSON with double quotes and no newlines
120
- clean_json = event.to_json(separators=(",", ":"), indent=None)
121
-
122
- return f"event: {event_type}\ndata: {clean_json}\n\n"
123
- except json.JSONDecodeError:
124
- clean_json = event.to_json(separators=(",", ":"), indent=None)
125
- return f"event: message\ndata: {clean_json}\n\n"
126
-
127
-
128
- class WebSocketManager:
129
- """Manages WebSocket connections for workflow runs"""
130
-
131
- active_connections: Dict[str, WebSocket] # {run_id: websocket}
132
- authenticated_connections: Dict[WebSocket, bool] # {websocket: is_authenticated}
133
-
134
- def __init__(
135
- self,
136
- active_connections: Optional[Dict[str, WebSocket]] = None,
137
- ):
138
- # Store active connections: {run_id: websocket}
139
- self.active_connections = active_connections or {}
140
- # Track authentication state for each websocket
141
- self.authenticated_connections = {}
142
-
143
- async def connect(self, websocket: WebSocket, requires_auth: bool = True):
144
- """Accept WebSocket connection"""
145
- await websocket.accept()
146
- logger.debug("WebSocket connected")
147
-
148
- # If auth is not required, mark as authenticated immediately
149
- self.authenticated_connections[websocket] = not requires_auth
150
-
151
- # Send connection confirmation with auth requirement info
152
- await websocket.send_text(
153
- json.dumps(
154
- {
155
- "event": "connected",
156
- "message": (
157
- "Connected to workflow events. Please authenticate to continue."
158
- if requires_auth
159
- else "Connected to workflow events. Authentication not required."
160
- ),
161
- "requires_auth": requires_auth,
162
- }
163
- )
164
- )
165
-
166
- async def authenticate_websocket(self, websocket: WebSocket):
167
- """Mark a WebSocket connection as authenticated"""
168
- self.authenticated_connections[websocket] = True
169
- logger.debug("WebSocket authenticated")
170
-
171
- # Send authentication confirmation
172
- await websocket.send_text(
173
- json.dumps(
174
- {
175
- "event": "authenticated",
176
- "message": "Authentication successful. You can now send commands.",
177
- }
178
- )
179
- )
180
-
181
- def is_authenticated(self, websocket: WebSocket) -> bool:
182
- """Check if a WebSocket connection is authenticated"""
183
- return self.authenticated_connections.get(websocket, False)
184
-
185
- async def register_workflow_websocket(self, run_id: str, websocket: WebSocket):
186
- """Register a workflow run with its WebSocket connection"""
187
- self.active_connections[run_id] = websocket
188
- logger.debug(f"Registered WebSocket for run_id: {run_id}")
189
-
190
- async def disconnect_by_run_id(self, run_id: str):
191
- """Remove WebSocket connection by run_id"""
192
- if run_id in self.active_connections:
193
- websocket = self.active_connections[run_id]
194
- del self.active_connections[run_id]
195
- # Clean up authentication state
196
- if websocket in self.authenticated_connections:
197
- del self.authenticated_connections[websocket]
198
- logger.debug(f"WebSocket disconnected for run_id: {run_id}")
199
-
200
- async def disconnect_websocket(self, websocket: WebSocket):
201
- """Remove WebSocket connection and clean up all associated state"""
202
- # Remove from authenticated connections
203
- if websocket in self.authenticated_connections:
204
- del self.authenticated_connections[websocket]
205
-
206
- # Remove from active connections
207
- runs_to_remove = [run_id for run_id, ws in self.active_connections.items() if ws == websocket]
208
- for run_id in runs_to_remove:
209
- del self.active_connections[run_id]
210
-
211
- logger.debug("WebSocket disconnected and cleaned up")
212
-
213
- async def get_websocket_for_run(self, run_id: str) -> Optional[WebSocket]:
214
- """Get WebSocket connection for a workflow run"""
215
- return self.active_connections.get(run_id)
216
-
217
-
218
- # Global manager instance
219
- websocket_manager = WebSocketManager(
220
- active_connections={},
221
- )
222
-
223
-
224
- async def agent_response_streamer(
225
- agent: Agent,
226
- message: str,
227
- session_id: Optional[str] = None,
228
- user_id: Optional[str] = None,
229
- images: Optional[List[Image]] = None,
230
- audio: Optional[List[Audio]] = None,
231
- videos: Optional[List[Video]] = None,
232
- files: Optional[List[FileMedia]] = None,
233
- **kwargs: Any,
234
- ) -> AsyncGenerator:
235
- try:
236
- run_response = agent.arun(
237
- input=message,
238
- session_id=session_id,
239
- user_id=user_id,
240
- images=images,
241
- audio=audio,
242
- videos=videos,
243
- files=files,
244
- stream=True,
245
- stream_intermediate_steps=True,
246
- **kwargs,
247
- )
248
- async for run_response_chunk in run_response:
249
- yield format_sse_event(run_response_chunk) # type: ignore
250
- except (InputCheckError, OutputCheckError) as e:
251
- error_response = RunErrorEvent(
252
- content=str(e),
253
- error_type=e.type,
254
- error_id=e.error_id,
255
- additional_data=e.additional_data,
256
- )
257
- yield format_sse_event(error_response)
258
- except Exception as e:
259
- import traceback
260
-
261
- traceback.print_exc(limit=3)
262
- error_response = RunErrorEvent(
263
- content=str(e),
264
- )
265
- yield format_sse_event(error_response)
266
-
267
-
268
- async def agent_continue_response_streamer(
269
- agent: Agent,
270
- run_id: Optional[str] = None,
271
- updated_tools: Optional[List] = None,
272
- session_id: Optional[str] = None,
273
- user_id: Optional[str] = None,
274
- ) -> AsyncGenerator:
275
- try:
276
- continue_response = agent.acontinue_run(
277
- run_id=run_id,
278
- updated_tools=updated_tools,
279
- session_id=session_id,
280
- user_id=user_id,
281
- stream=True,
282
- stream_intermediate_steps=True,
283
- )
284
- async for run_response_chunk in continue_response:
285
- yield format_sse_event(run_response_chunk) # type: ignore
286
- except (InputCheckError, OutputCheckError) as e:
287
- error_response = RunErrorEvent(
288
- content=str(e),
289
- error_type=e.type,
290
- error_id=e.error_id,
291
- additional_data=e.additional_data,
292
- )
293
- yield format_sse_event(error_response)
294
-
295
- except Exception as e:
296
- import traceback
297
-
298
- traceback.print_exc(limit=3)
299
- error_response = RunErrorEvent(
300
- content=str(e),
301
- error_type=e.type if hasattr(e, "type") else None,
302
- error_id=e.error_id if hasattr(e, "error_id") else None,
303
- )
304
- yield format_sse_event(error_response)
305
- return
306
-
307
-
308
- async def team_response_streamer(
309
- team: Team,
310
- message: str,
311
- session_id: Optional[str] = None,
312
- user_id: Optional[str] = None,
313
- images: Optional[List[Image]] = None,
314
- audio: Optional[List[Audio]] = None,
315
- videos: Optional[List[Video]] = None,
316
- files: Optional[List[FileMedia]] = None,
317
- **kwargs: Any,
318
- ) -> AsyncGenerator:
319
- """Run the given team asynchronously and yield its response"""
320
- try:
321
- run_response = team.arun(
322
- input=message,
323
- session_id=session_id,
324
- user_id=user_id,
325
- images=images,
326
- audio=audio,
327
- videos=videos,
328
- files=files,
329
- stream=True,
330
- stream_intermediate_steps=True,
331
- **kwargs,
332
- )
333
- async for run_response_chunk in run_response:
334
- yield format_sse_event(run_response_chunk) # type: ignore
335
- except (InputCheckError, OutputCheckError) as e:
336
- error_response = TeamRunErrorEvent(
337
- content=str(e),
338
- error_type=e.type,
339
- error_id=e.error_id,
340
- additional_data=e.additional_data,
341
- )
342
- yield format_sse_event(error_response)
343
-
344
- except Exception as e:
345
- import traceback
346
-
347
- traceback.print_exc()
348
- error_response = TeamRunErrorEvent(
349
- content=str(e),
350
- error_type=e.type if hasattr(e, "type") else None,
351
- error_id=e.error_id if hasattr(e, "error_id") else None,
352
- )
353
- yield format_sse_event(error_response)
354
- return
355
-
356
-
357
- async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os: "AgentOS"):
358
- """Handle workflow execution directly via WebSocket"""
359
- try:
360
- workflow_id = message.get("workflow_id")
361
- session_id = message.get("session_id")
362
- user_message = message.get("message", "")
363
- user_id = message.get("user_id")
364
-
365
- if not workflow_id:
366
- await websocket.send_text(json.dumps({"event": "error", "error": "workflow_id is required"}))
367
- return
368
-
369
- # Get workflow from OS
370
- workflow = get_workflow_by_id(workflow_id, os.workflows)
371
- if not workflow:
372
- await websocket.send_text(json.dumps({"event": "error", "error": f"Workflow {workflow_id} not found"}))
373
- return
374
-
375
- # Generate session_id if not provided
376
- # Use workflow's default session_id if not provided in message
377
- if not session_id:
378
- if workflow.session_id:
379
- session_id = workflow.session_id
380
- else:
381
- session_id = str(uuid4())
382
-
383
- # Execute workflow in background with streaming
384
- workflow_result = await workflow.arun(
385
- input=user_message,
386
- session_id=session_id,
387
- user_id=user_id,
388
- stream=True,
389
- stream_intermediate_steps=True,
390
- background=True,
391
- websocket=websocket,
392
- )
393
-
394
- workflow_run_output = cast(WorkflowRunOutput, workflow_result)
395
-
396
- await websocket_manager.register_workflow_websocket(workflow_run_output.run_id, websocket) # type: ignore
397
-
398
- except (InputCheckError, OutputCheckError) as e:
399
- await websocket.send_text(
400
- json.dumps(
401
- {
402
- "event": "error",
403
- "error": str(e),
404
- "error_type": e.type,
405
- "error_id": e.error_id,
406
- "additional_data": e.additional_data,
407
- }
408
- )
409
- )
410
- except Exception as e:
411
- logger.error(f"Error executing workflow via WebSocket: {e}")
412
- error_payload = {
413
- "event": "error",
414
- "error": str(e),
415
- "error_type": e.type if hasattr(e, "type") else None,
416
- "error_id": e.error_id if hasattr(e, "error_id") else None,
417
- }
418
- error_payload = {k: v for k, v in error_payload.items() if v is not None}
419
- await websocket.send_text(json.dumps(error_payload))
420
-
421
-
422
- async def workflow_response_streamer(
423
- workflow: Workflow,
424
- input: Optional[Union[str, Dict[str, Any], List[Any], BaseModel]] = None,
425
- session_id: Optional[str] = None,
426
- user_id: Optional[str] = None,
427
- **kwargs: Any,
428
- ) -> AsyncGenerator:
429
- try:
430
- run_response = await workflow.arun(
431
- input=input,
432
- session_id=session_id,
433
- user_id=user_id,
434
- stream=True,
435
- stream_intermediate_steps=True,
436
- **kwargs,
437
- )
438
-
439
- async for run_response_chunk in run_response:
440
- yield format_sse_event(run_response_chunk) # type: ignore
441
-
442
- except (InputCheckError, OutputCheckError) as e:
443
- error_response = WorkflowErrorEvent(
444
- error=str(e),
445
- error_type=e.type,
446
- error_id=e.error_id,
447
- additional_data=e.additional_data,
448
- )
449
- yield format_sse_event(error_response)
450
-
451
- except Exception as e:
452
- import traceback
453
-
454
- traceback.print_exc()
455
- error_response = WorkflowErrorEvent(
456
- error=str(e),
457
- error_type=e.type if hasattr(e, "type") else None,
458
- error_id=e.error_id if hasattr(e, "error_id") else None,
459
- )
460
- yield format_sse_event(error_response)
461
- return
462
-
463
-
464
- def get_websocket_router(
465
- os: "AgentOS",
466
- settings: AgnoAPISettings = AgnoAPISettings(),
467
- ) -> APIRouter:
468
- """
469
- Create WebSocket router without HTTP authentication dependencies.
470
- WebSocket endpoints handle authentication internally via message-based auth.
471
- """
472
- ws_router = APIRouter()
473
-
474
- @ws_router.websocket(
475
- "/workflows/ws",
476
- name="workflow_websocket",
477
- )
478
- async def workflow_websocket_endpoint(websocket: WebSocket):
479
- """WebSocket endpoint for receiving real-time workflow events"""
480
- requires_auth = bool(settings.os_security_key)
481
- await websocket_manager.connect(websocket, requires_auth=requires_auth)
482
-
483
- try:
484
- while True:
485
- data = await websocket.receive_text()
486
- message = json.loads(data)
487
- action = message.get("action")
488
-
489
- # Handle authentication first
490
- if action == "authenticate":
491
- token = message.get("token")
492
- if not token:
493
- await websocket.send_text(json.dumps({"event": "auth_error", "error": "Token is required"}))
494
- continue
495
-
496
- if validate_websocket_token(token, settings):
497
- await websocket_manager.authenticate_websocket(websocket)
498
- else:
499
- await websocket.send_text(json.dumps({"event": "auth_error", "error": "Invalid token"}))
500
- continue
501
-
502
- # Check authentication for all other actions (only when required)
503
- elif requires_auth and not websocket_manager.is_authenticated(websocket):
504
- await websocket.send_text(
505
- json.dumps(
506
- {
507
- "event": "auth_required",
508
- "error": "Authentication required. Send authenticate action with valid token.",
509
- }
510
- )
511
- )
512
- continue
513
-
514
- # Handle authenticated actions
515
- elif action == "ping":
516
- await websocket.send_text(json.dumps({"event": "pong"}))
517
-
518
- elif action == "start-workflow":
519
- # Handle workflow execution directly via WebSocket
520
- await handle_workflow_via_websocket(websocket, message, os)
521
-
522
- else:
523
- await websocket.send_text(json.dumps({"event": "error", "error": f"Unknown action: {action}"}))
524
-
525
- except Exception as e:
526
- if "1012" not in str(e) and "1001" not in str(e):
527
- logger.error(f"WebSocket error: {e}")
528
- finally:
529
- # Clean up the websocket connection
530
- await websocket_manager.disconnect_websocket(websocket)
531
-
532
- return ws_router
533
-
534
-
535
38
  def get_base_router(
536
39
  os: "AgentOS",
537
40
  settings: AgnoAPISettings = AgnoAPISettings(),
@@ -644,13 +147,14 @@ def get_base_router(
644
147
  os_id=os.id or "Unnamed OS",
645
148
  description=os.description,
646
149
  available_models=os.config.available_models if os.config else [],
647
- databases=list({db.id for db in chain(os.dbs.values(), os.knowledge_dbs.values())}),
150
+ databases=list({db.id for db_id, dbs in os.dbs.items() for db in dbs}),
648
151
  chat=os.config.chat if os.config else None,
649
152
  session=os._get_session_config(),
650
153
  memory=os._get_memory_config(),
651
154
  knowledge=os._get_knowledge_config(),
652
155
  evals=os._get_evals_config(),
653
156
  metrics=os._get_metrics_config(),
157
+ traces=os._get_traces_config(),
654
158
  agents=[AgentSummaryResponse.from_agent(agent) for agent in os.agents] if os.agents else [],
655
159
  teams=[TeamSummaryResponse.from_team(team) for team in os.teams] if os.teams else [],
656
160
  workflows=[WorkflowSummaryResponse.from_workflow(w) for w in os.workflows] if os.workflows else [],
@@ -703,994 +207,52 @@ def get_base_router(
703
207
 
704
208
  return list(unique_models.values())
705
209
 
706
- # -- Agent routes ---
707
-
708
- @router.post(
709
- "/agents/{agent_id}/runs",
710
- tags=["Agents"],
711
- operation_id="create_agent_run",
712
- response_model_exclude_none=True,
713
- summary="Create Agent Run",
714
- description=(
715
- "Execute an agent with a message and optional media files. Supports both streaming and non-streaming responses.\n\n"
716
- "**Features:**\n"
717
- "- Text message input with optional session management\n"
718
- "- Multi-media support: images (PNG, JPEG, WebP), audio (WAV, MP3), video (MP4, WebM, etc.)\n"
719
- "- Document processing: PDF, CSV, DOCX, TXT, JSON\n"
720
- "- Real-time streaming responses with Server-Sent Events (SSE)\n"
721
- "- User and session context preservation\n\n"
722
- "**Streaming Response:**\n"
723
- "When `stream=true`, returns SSE events with `event` and `data` fields."
724
- ),
725
- responses={
726
- 200: {
727
- "description": "Agent run executed successfully",
728
- "content": {
729
- "text/event-stream": {
730
- "examples": {
731
- "event_stream": {
732
- "summary": "Example event stream response",
733
- "value": 'event: RunStarted\ndata: {"content": "Hello!", "run_id": "123..."}\n\n',
734
- }
735
- }
736
- },
737
- },
738
- },
739
- 400: {"description": "Invalid request or unsupported file type", "model": BadRequestResponse},
740
- 404: {"description": "Agent not found", "model": NotFoundResponse},
741
- },
742
- )
743
- async def create_agent_run(
744
- agent_id: str,
745
- request: Request,
746
- message: str = Form(...),
747
- stream: bool = Form(False),
748
- session_id: Optional[str] = Form(None),
749
- user_id: Optional[str] = Form(None),
750
- files: Optional[List[UploadFile]] = File(None),
751
- ):
752
- kwargs = await _get_request_kwargs(request, create_agent_run)
753
-
754
- if hasattr(request.state, "user_id"):
755
- if user_id:
756
- log_warning("User ID parameter passed in both request state and kwargs, using request state")
757
- user_id = request.state.user_id
758
- if hasattr(request.state, "session_id"):
759
- if session_id:
760
- log_warning("Session ID parameter passed in both request state and kwargs, using request state")
761
- session_id = request.state.session_id
762
- if hasattr(request.state, "session_state"):
763
- session_state = request.state.session_state
764
- if "session_state" in kwargs:
765
- log_warning("Session state parameter passed in both request state and kwargs, using request state")
766
- kwargs["session_state"] = session_state
767
- if hasattr(request.state, "dependencies"):
768
- dependencies = request.state.dependencies
769
- if "dependencies" in kwargs:
770
- log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
771
- kwargs["dependencies"] = dependencies
772
-
773
- agent = get_agent_by_id(agent_id, os.agents)
774
- if agent is None:
775
- raise HTTPException(status_code=404, detail="Agent not found")
776
-
777
- if session_id is None or session_id == "":
778
- log_debug("Creating new session")
779
- session_id = str(uuid4())
780
-
781
- base64_images: List[Image] = []
782
- base64_audios: List[Audio] = []
783
- base64_videos: List[Video] = []
784
- input_files: List[FileMedia] = []
785
-
786
- if files:
787
- for file in files:
788
- if file.content_type in [
789
- "image/png",
790
- "image/jpeg",
791
- "image/jpg",
792
- "image/gif",
793
- "image/webp",
794
- "image/bmp",
795
- "image/tiff",
796
- "image/tif",
797
- "image/avif",
798
- ]:
799
- try:
800
- base64_image = process_image(file)
801
- base64_images.append(base64_image)
802
- except Exception as e:
803
- log_error(f"Error processing image {file.filename}: {e}")
804
- continue
805
- elif file.content_type in [
806
- "audio/wav",
807
- "audio/wave",
808
- "audio/mp3",
809
- "audio/mpeg",
810
- "audio/ogg",
811
- "audio/mp4",
812
- "audio/m4a",
813
- "audio/aac",
814
- "audio/flac",
815
- ]:
816
- try:
817
- audio = process_audio(file)
818
- base64_audios.append(audio)
819
- except Exception as e:
820
- log_error(f"Error processing audio {file.filename} with content type {file.content_type}: {e}")
821
- continue
822
- elif file.content_type in [
823
- "video/x-flv",
824
- "video/quicktime",
825
- "video/mpeg",
826
- "video/mpegs",
827
- "video/mpgs",
828
- "video/mpg",
829
- "video/mpg",
830
- "video/mp4",
831
- "video/webm",
832
- "video/wmv",
833
- "video/3gpp",
834
- ]:
835
- try:
836
- base64_video = process_video(file)
837
- base64_videos.append(base64_video)
838
- except Exception as e:
839
- log_error(f"Error processing video {file.filename}: {e}")
840
- continue
841
- elif file.content_type in [
842
- "application/pdf",
843
- "application/json",
844
- "application/x-javascript",
845
- "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
846
- "text/javascript",
847
- "application/x-python",
848
- "text/x-python",
849
- "text/plain",
850
- "text/html",
851
- "text/css",
852
- "text/md",
853
- "text/csv",
854
- "text/xml",
855
- "text/rtf",
856
- ]:
857
- # Process document files
858
- try:
859
- input_file = process_document(file)
860
- if input_file is not None:
861
- input_files.append(input_file)
862
- except Exception as e:
863
- log_error(f"Error processing file {file.filename}: {e}")
864
- continue
865
- else:
866
- raise HTTPException(status_code=400, detail="Unsupported file type")
867
-
868
- if stream:
869
- return StreamingResponse(
870
- agent_response_streamer(
871
- agent,
872
- message,
873
- session_id=session_id,
874
- user_id=user_id,
875
- images=base64_images if base64_images else None,
876
- audio=base64_audios if base64_audios else None,
877
- videos=base64_videos if base64_videos else None,
878
- files=input_files if input_files else None,
879
- **kwargs,
880
- ),
881
- media_type="text/event-stream",
882
- )
883
- else:
884
- try:
885
- run_response = cast(
886
- RunOutput,
887
- await agent.arun(
888
- input=message,
889
- session_id=session_id,
890
- user_id=user_id,
891
- images=base64_images if base64_images else None,
892
- audio=base64_audios if base64_audios else None,
893
- videos=base64_videos if base64_videos else None,
894
- files=input_files if input_files else None,
895
- stream=False,
896
- **kwargs,
897
- ),
898
- )
899
- return run_response.to_dict()
900
-
901
- except InputCheckError as e:
902
- raise HTTPException(status_code=400, detail=str(e))
903
-
904
- @router.post(
905
- "/agents/{agent_id}/runs/{run_id}/cancel",
906
- tags=["Agents"],
907
- operation_id="cancel_agent_run",
908
- response_model_exclude_none=True,
909
- summary="Cancel Agent Run",
910
- description=(
911
- "Cancel a currently executing agent run. This will attempt to stop the agent's execution gracefully.\n\n"
912
- "**Note:** Cancellation may not be immediate for all operations."
913
- ),
914
- responses={
915
- 200: {},
916
- 404: {"description": "Agent not found", "model": NotFoundResponse},
917
- 500: {"description": "Failed to cancel run", "model": InternalServerErrorResponse},
918
- },
919
- )
920
- async def cancel_agent_run(
921
- agent_id: str,
922
- run_id: str,
923
- ):
924
- agent = get_agent_by_id(agent_id, os.agents)
925
- if agent is None:
926
- raise HTTPException(status_code=404, detail="Agent not found")
927
-
928
- if not agent.cancel_run(run_id=run_id):
929
- raise HTTPException(status_code=500, detail="Failed to cancel run")
930
-
931
- return JSONResponse(content={}, status_code=200)
932
-
933
- @router.post(
934
- "/agents/{agent_id}/runs/{run_id}/continue",
935
- tags=["Agents"],
936
- operation_id="continue_agent_run",
937
- response_model_exclude_none=True,
938
- summary="Continue Agent Run",
939
- description=(
940
- "Continue a paused or incomplete agent run with updated tool results.\n\n"
941
- "**Use Cases:**\n"
942
- "- Resume execution after tool approval/rejection\n"
943
- "- Provide manual tool execution results\n\n"
944
- "**Tools Parameter:**\n"
945
- "JSON string containing array of tool execution objects with results."
946
- ),
947
- responses={
948
- 200: {
949
- "description": "Agent run continued successfully",
950
- "content": {
951
- "text/event-stream": {
952
- "example": 'event: RunContent\ndata: {"created_at": 1757348314, "run_id": "123..."}\n\n'
953
- },
954
- },
955
- },
956
- 400: {"description": "Invalid JSON in tools field or invalid tool structure", "model": BadRequestResponse},
957
- 404: {"description": "Agent not found", "model": NotFoundResponse},
958
- },
959
- )
960
- async def continue_agent_run(
961
- agent_id: str,
962
- run_id: str,
963
- request: Request,
964
- tools: str = Form(...), # JSON string of tools
965
- session_id: Optional[str] = Form(None),
966
- user_id: Optional[str] = Form(None),
967
- stream: bool = Form(True),
968
- ):
969
- if hasattr(request.state, "user_id"):
970
- user_id = request.state.user_id
971
- if hasattr(request.state, "session_id"):
972
- session_id = request.state.session_id
973
-
974
- # Parse the JSON string manually
975
- try:
976
- tools_data = json.loads(tools) if tools else None
977
- except json.JSONDecodeError:
978
- raise HTTPException(status_code=400, detail="Invalid JSON in tools field")
979
-
980
- agent = get_agent_by_id(agent_id, os.agents)
981
- if agent is None:
982
- raise HTTPException(status_code=404, detail="Agent not found")
983
-
984
- if session_id is None or session_id == "":
985
- log_warning(
986
- "Continuing run without session_id. This might lead to unexpected behavior if session context is important."
987
- )
988
-
989
- # Convert tools dict to ToolExecution objects if provided
990
- updated_tools = None
991
- if tools_data:
992
- try:
993
- from agno.models.response import ToolExecution
994
-
995
- updated_tools = [ToolExecution.from_dict(tool) for tool in tools_data]
996
- except Exception as e:
997
- raise HTTPException(status_code=400, detail=f"Invalid structure or content for tools: {str(e)}")
998
-
999
- if stream:
1000
- return StreamingResponse(
1001
- agent_continue_response_streamer(
1002
- agent,
1003
- run_id=run_id, # run_id from path
1004
- updated_tools=updated_tools,
1005
- session_id=session_id,
1006
- user_id=user_id,
1007
- ),
1008
- media_type="text/event-stream",
1009
- )
1010
- else:
1011
- try:
1012
- run_response_obj = cast(
1013
- RunOutput,
1014
- await agent.acontinue_run(
1015
- run_id=run_id, # run_id from path
1016
- updated_tools=updated_tools,
1017
- session_id=session_id,
1018
- user_id=user_id,
1019
- stream=False,
1020
- ),
1021
- )
1022
- return run_response_obj.to_dict()
1023
-
1024
- except InputCheckError as e:
1025
- raise HTTPException(status_code=400, detail=str(e))
1026
-
1027
- @router.get(
1028
- "/agents",
1029
- response_model=List[AgentResponse],
1030
- response_model_exclude_none=True,
1031
- tags=["Agents"],
1032
- operation_id="get_agents",
1033
- summary="List All Agents",
1034
- description=(
1035
- "Retrieve a comprehensive list of all agents configured in this OS instance.\n\n"
1036
- "**Returns:**\n"
1037
- "- Agent metadata (ID, name, description)\n"
1038
- "- Model configuration and capabilities\n"
1039
- "- Available tools and their configurations\n"
1040
- "- Session, knowledge, memory, and reasoning settings\n"
1041
- "- Only meaningful (non-default) configurations are included"
1042
- ),
1043
- responses={
1044
- 200: {
1045
- "description": "List of agents retrieved successfully",
1046
- "content": {
1047
- "application/json": {
1048
- "example": [
1049
- {
1050
- "id": "main-agent",
1051
- "name": "Main Agent",
1052
- "db_id": "c6bf0644-feb8-4930-a305-380dae5ad6aa",
1053
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI"},
1054
- "tools": None,
1055
- "sessions": {"session_table": "agno_sessions"},
1056
- "knowledge": {"knowledge_table": "main_knowledge"},
1057
- "system_message": {"markdown": True, "add_datetime_to_context": True},
1058
- }
1059
- ]
1060
- }
1061
- },
1062
- }
1063
- },
1064
- )
1065
- async def get_agents() -> List[AgentResponse]:
1066
- """Return the list of all Agents present in the contextual OS"""
1067
- if os.agents is None:
1068
- return []
1069
-
1070
- agents = []
1071
- for agent in os.agents:
1072
- agents.append(AgentResponse.from_agent(agent=agent))
1073
-
1074
- return agents
1075
-
1076
- @router.get(
1077
- "/agents/{agent_id}",
1078
- response_model=AgentResponse,
1079
- response_model_exclude_none=True,
1080
- tags=["Agents"],
1081
- operation_id="get_agent",
1082
- summary="Get Agent Details",
1083
- description=(
1084
- "Retrieve detailed configuration and capabilities of a specific agent.\n\n"
1085
- "**Returns comprehensive agent information including:**\n"
1086
- "- Model configuration and provider details\n"
1087
- "- Complete tool inventory and configurations\n"
1088
- "- Session management settings\n"
1089
- "- Knowledge base and memory configurations\n"
1090
- "- Reasoning capabilities and settings\n"
1091
- "- System prompts and response formatting options"
1092
- ),
1093
- responses={
1094
- 200: {
1095
- "description": "Agent details retrieved successfully",
1096
- "content": {
1097
- "application/json": {
1098
- "example": {
1099
- "id": "main-agent",
1100
- "name": "Main Agent",
1101
- "db_id": "9e064c70-6821-4840-a333-ce6230908a70",
1102
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI"},
1103
- "tools": None,
1104
- "sessions": {"session_table": "agno_sessions"},
1105
- "knowledge": {"knowledge_table": "main_knowledge"},
1106
- "system_message": {"markdown": True, "add_datetime_to_context": True},
1107
- }
1108
- }
1109
- },
1110
- },
1111
- 404: {"description": "Agent not found", "model": NotFoundResponse},
1112
- },
1113
- )
1114
- async def get_agent(agent_id: str) -> AgentResponse:
1115
- agent = get_agent_by_id(agent_id, os.agents)
1116
- if agent is None:
1117
- raise HTTPException(status_code=404, detail="Agent not found")
1118
-
1119
- return AgentResponse.from_agent(agent)
1120
-
1121
- # -- Team routes ---
1122
-
1123
- @router.post(
1124
- "/teams/{team_id}/runs",
1125
- tags=["Teams"],
1126
- operation_id="create_team_run",
1127
- response_model_exclude_none=True,
1128
- summary="Create Team Run",
1129
- description=(
1130
- "Execute a team collaboration with multiple agents working together on a task.\n\n"
1131
- "**Features:**\n"
1132
- "- Text message input with optional session management\n"
1133
- "- Multi-media support: images (PNG, JPEG, WebP), audio (WAV, MP3), video (MP4, WebM, etc.)\n"
1134
- "- Document processing: PDF, CSV, DOCX, TXT, JSON\n"
1135
- "- Real-time streaming responses with Server-Sent Events (SSE)\n"
1136
- "- User and session context preservation\n\n"
1137
- "**Streaming Response:**\n"
1138
- "When `stream=true`, returns SSE events with `event` and `data` fields."
1139
- ),
1140
- responses={
1141
- 200: {
1142
- "description": "Team run executed successfully",
1143
- "content": {
1144
- "text/event-stream": {
1145
- "example": 'event: RunStarted\ndata: {"content": "Hello!", "run_id": "123..."}\n\n'
1146
- },
1147
- },
1148
- },
1149
- 400: {"description": "Invalid request or unsupported file type", "model": BadRequestResponse},
1150
- 404: {"description": "Team not found", "model": NotFoundResponse},
1151
- },
1152
- )
1153
- async def create_team_run(
1154
- team_id: str,
1155
- request: Request,
1156
- message: str = Form(...),
1157
- stream: bool = Form(True),
1158
- monitor: bool = Form(True),
1159
- session_id: Optional[str] = Form(None),
1160
- user_id: Optional[str] = Form(None),
1161
- files: Optional[List[UploadFile]] = File(None),
1162
- ):
1163
- kwargs = await _get_request_kwargs(request, create_team_run)
1164
-
1165
- if hasattr(request.state, "user_id"):
1166
- if user_id:
1167
- log_warning("User ID parameter passed in both request state and kwargs, using request state")
1168
- user_id = request.state.user_id
1169
- if hasattr(request.state, "session_id"):
1170
- if session_id:
1171
- log_warning("Session ID parameter passed in both request state and kwargs, using request state")
1172
- session_id = request.state.session_id
1173
- if hasattr(request.state, "session_state"):
1174
- session_state = request.state.session_state
1175
- if "session_state" in kwargs:
1176
- log_warning("Session state parameter passed in both request state and kwargs, using request state")
1177
- kwargs["session_state"] = session_state
1178
- if hasattr(request.state, "dependencies"):
1179
- dependencies = request.state.dependencies
1180
- if "dependencies" in kwargs:
1181
- log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
1182
- kwargs["dependencies"] = dependencies
1183
-
1184
- logger.debug(f"Creating team run: {message=} {session_id=} {monitor=} {user_id=} {team_id=} {files=} {kwargs=}")
1185
-
1186
- team = get_team_by_id(team_id, os.teams)
1187
- if team is None:
1188
- raise HTTPException(status_code=404, detail="Team not found")
1189
-
1190
- if session_id is not None and session_id != "":
1191
- logger.debug(f"Continuing session: {session_id}")
1192
- else:
1193
- logger.debug("Creating new session")
1194
- session_id = str(uuid4())
1195
-
1196
- base64_images: List[Image] = []
1197
- base64_audios: List[Audio] = []
1198
- base64_videos: List[Video] = []
1199
- document_files: List[FileMedia] = []
1200
-
1201
- if files:
1202
- for file in files:
1203
- if file.content_type in ["image/png", "image/jpeg", "image/jpg", "image/webp"]:
1204
- try:
1205
- base64_image = process_image(file)
1206
- base64_images.append(base64_image)
1207
- except Exception as e:
1208
- logger.error(f"Error processing image {file.filename}: {e}")
1209
- continue
1210
- elif file.content_type in ["audio/wav", "audio/mp3", "audio/mpeg"]:
1211
- try:
1212
- base64_audio = process_audio(file)
1213
- base64_audios.append(base64_audio)
1214
- except Exception as e:
1215
- logger.error(f"Error processing audio {file.filename}: {e}")
1216
- continue
1217
- elif file.content_type in [
1218
- "video/x-flv",
1219
- "video/quicktime",
1220
- "video/mpeg",
1221
- "video/mpegs",
1222
- "video/mpgs",
1223
- "video/mpg",
1224
- "video/mpg",
1225
- "video/mp4",
1226
- "video/webm",
1227
- "video/wmv",
1228
- "video/3gpp",
1229
- ]:
1230
- try:
1231
- base64_video = process_video(file)
1232
- base64_videos.append(base64_video)
1233
- except Exception as e:
1234
- logger.error(f"Error processing video {file.filename}: {e}")
1235
- continue
1236
- elif file.content_type in [
1237
- "application/pdf",
1238
- "text/csv",
1239
- "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
1240
- "text/plain",
1241
- "application/json",
1242
- ]:
1243
- document_file = process_document(file)
1244
- if document_file is not None:
1245
- document_files.append(document_file)
1246
- else:
1247
- raise HTTPException(status_code=400, detail="Unsupported file type")
1248
-
1249
- if stream:
1250
- return StreamingResponse(
1251
- team_response_streamer(
1252
- team,
1253
- message,
1254
- session_id=session_id,
1255
- user_id=user_id,
1256
- images=base64_images if base64_images else None,
1257
- audio=base64_audios if base64_audios else None,
1258
- videos=base64_videos if base64_videos else None,
1259
- files=document_files if document_files else None,
1260
- **kwargs,
1261
- ),
1262
- media_type="text/event-stream",
1263
- )
1264
- else:
1265
- try:
1266
- run_response = await team.arun(
1267
- input=message,
1268
- session_id=session_id,
1269
- user_id=user_id,
1270
- images=base64_images if base64_images else None,
1271
- audio=base64_audios if base64_audios else None,
1272
- videos=base64_videos if base64_videos else None,
1273
- files=document_files if document_files else None,
1274
- stream=False,
1275
- **kwargs,
1276
- )
1277
- return run_response.to_dict()
1278
-
1279
- except InputCheckError as e:
1280
- raise HTTPException(status_code=400, detail=str(e))
1281
-
210
+ # -- Database Migration routes ---
1282
211
  @router.post(
1283
- "/teams/{team_id}/runs/{run_id}/cancel",
1284
- tags=["Teams"],
1285
- operation_id="cancel_team_run",
1286
- response_model_exclude_none=True,
1287
- summary="Cancel Team Run",
212
+ "/databases/{db_id}/migrate",
213
+ tags=["Database"],
214
+ operation_id="migrate_database",
215
+ summary="Migrate Database",
1288
216
  description=(
1289
- "Cancel a currently executing team run. This will attempt to stop the team's execution gracefully.\n\n"
1290
- "**Note:** Cancellation may not be immediate for all operations."
1291
- ),
1292
- responses={
1293
- 200: {},
1294
- 404: {"description": "Team not found", "model": NotFoundResponse},
1295
- 500: {"description": "Failed to cancel team run", "model": InternalServerErrorResponse},
1296
- },
1297
- )
1298
- async def cancel_team_run(
1299
- team_id: str,
1300
- run_id: str,
1301
- ):
1302
- team = get_team_by_id(team_id, os.teams)
1303
- if team is None:
1304
- raise HTTPException(status_code=404, detail="Team not found")
1305
-
1306
- if not team.cancel_run(run_id=run_id):
1307
- raise HTTPException(status_code=500, detail="Failed to cancel run")
1308
-
1309
- return JSONResponse(content={}, status_code=200)
1310
-
1311
- @router.get(
1312
- "/teams",
1313
- response_model=List[TeamResponse],
1314
- response_model_exclude_none=True,
1315
- tags=["Teams"],
1316
- operation_id="get_teams",
1317
- summary="List All Teams",
1318
- description=(
1319
- "Retrieve a comprehensive list of all teams configured in this OS instance.\n\n"
1320
- "**Returns team information including:**\n"
1321
- "- Team metadata (ID, name, description, execution mode)\n"
1322
- "- Model configuration for team coordination\n"
1323
- "- Team member roster with roles and capabilities\n"
1324
- "- Knowledge sharing and memory configurations"
217
+ "Migrate the given database schema to the given target version. "
218
+ "If a target version is not provided, the database will be migrated to the latest version. "
1325
219
  ),
1326
220
  responses={
1327
221
  200: {
1328
- "description": "List of teams retrieved successfully",
222
+ "description": "Database migrated successfully",
1329
223
  "content": {
1330
224
  "application/json": {
1331
- "example": [
1332
- {
1333
- "team_id": "basic-team",
1334
- "name": "Basic Team",
1335
- "mode": "coordinate",
1336
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI"},
1337
- "tools": [
1338
- {
1339
- "name": "transfer_task_to_member",
1340
- "description": "Use this function to transfer a task to the selected team member.\nYou must provide a clear and concise description of the task the member should achieve AND the expected output.",
1341
- "parameters": {
1342
- "type": "object",
1343
- "properties": {
1344
- "member_id": {
1345
- "type": "string",
1346
- "description": "(str) The ID of the member to transfer the task to. Use only the ID of the member, not the ID of the team followed by the ID of the member.",
1347
- },
1348
- "task_description": {
1349
- "type": "string",
1350
- "description": "(str) A clear and concise description of the task the member should achieve.",
1351
- },
1352
- "expected_output": {
1353
- "type": "string",
1354
- "description": "(str) The expected output from the member (optional).",
1355
- },
1356
- },
1357
- "additionalProperties": False,
1358
- "required": ["member_id", "task_description"],
1359
- },
1360
- }
1361
- ],
1362
- "members": [
1363
- {
1364
- "agent_id": "basic-agent",
1365
- "name": "Basic Agent",
1366
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI gpt-4o"},
1367
- "memory": {
1368
- "app_name": "Memory",
1369
- "app_url": None,
1370
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI"},
1371
- },
1372
- "session_table": "agno_sessions",
1373
- "memory_table": "agno_memories",
1374
- }
1375
- ],
1376
- "enable_agentic_context": False,
1377
- "memory": {
1378
- "app_name": "agno_memories",
1379
- "app_url": "/memory/1",
1380
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI"},
1381
- },
1382
- "async_mode": False,
1383
- "session_table": "agno_sessions",
1384
- "memory_table": "agno_memories",
1385
- }
1386
- ]
1387
- }
1388
- },
1389
- }
1390
- },
1391
- )
1392
- async def get_teams() -> List[TeamResponse]:
1393
- """Return the list of all Teams present in the contextual OS"""
1394
- if os.teams is None:
1395
- return []
1396
-
1397
- teams = []
1398
- for team in os.teams:
1399
- teams.append(TeamResponse.from_team(team=team))
1400
-
1401
- return teams
1402
-
1403
- @router.get(
1404
- "/teams/{team_id}",
1405
- response_model=TeamResponse,
1406
- response_model_exclude_none=True,
1407
- tags=["Teams"],
1408
- operation_id="get_team",
1409
- summary="Get Team Details",
1410
- description=("Retrieve detailed configuration and member information for a specific team."),
1411
- responses={
1412
- 200: {
1413
- "description": "Team details retrieved successfully",
1414
- "content": {
1415
- "application/json": {
1416
- "example": {
1417
- "team_id": "basic-team",
1418
- "name": "Basic Team",
1419
- "description": None,
1420
- "mode": "coordinate",
1421
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI"},
1422
- "tools": [
1423
- {
1424
- "name": "transfer_task_to_member",
1425
- "description": "Use this function to transfer a task to the selected team member.\nYou must provide a clear and concise description of the task the member should achieve AND the expected output.",
1426
- "parameters": {
1427
- "type": "object",
1428
- "properties": {
1429
- "member_id": {
1430
- "type": "string",
1431
- "description": "(str) The ID of the member to transfer the task to. Use only the ID of the member, not the ID of the team followed by the ID of the member.",
1432
- },
1433
- "task_description": {
1434
- "type": "string",
1435
- "description": "(str) A clear and concise description of the task the member should achieve.",
1436
- },
1437
- "expected_output": {
1438
- "type": "string",
1439
- "description": "(str) The expected output from the member (optional).",
1440
- },
1441
- },
1442
- "additionalProperties": False,
1443
- "required": ["member_id", "task_description"],
1444
- },
1445
- }
1446
- ],
1447
- "instructions": None,
1448
- "members": [
1449
- {
1450
- "agent_id": "basic-agent",
1451
- "name": "Basic Agent",
1452
- "description": None,
1453
- "instructions": None,
1454
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI gpt-4o"},
1455
- "tools": None,
1456
- "memory": {
1457
- "app_name": "Memory",
1458
- "app_url": None,
1459
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI"},
1460
- },
1461
- "knowledge": None,
1462
- "session_table": "agno_sessions",
1463
- "memory_table": "agno_memories",
1464
- "knowledge_table": None,
1465
- }
1466
- ],
1467
- "expected_output": None,
1468
- "dependencies": None,
1469
- "enable_agentic_context": False,
1470
- "memory": {
1471
- "app_name": "Memory",
1472
- "app_url": None,
1473
- "model": {"name": "OpenAIChat", "model": "gpt-4o", "provider": "OpenAI"},
1474
- },
1475
- "knowledge": None,
1476
- "async_mode": False,
1477
- "session_table": "agno_sessions",
1478
- "memory_table": "agno_memories",
1479
- "knowledge_table": None,
1480
- }
225
+ "example": {"message": "Database migrated successfully to version 3.0.0"},
1481
226
  }
1482
227
  },
1483
228
  },
1484
- 404: {"description": "Team not found", "model": NotFoundResponse},
229
+ 404: {"description": "Database not found", "model": NotFoundResponse},
230
+ 500: {"description": "Failed to migrate database", "model": InternalServerErrorResponse},
1485
231
  },
1486
232
  )
1487
- async def get_team(team_id: str) -> TeamResponse:
1488
- team = get_team_by_id(team_id, os.teams)
1489
- if team is None:
1490
- raise HTTPException(status_code=404, detail="Team not found")
1491
-
1492
- return TeamResponse.from_team(team)
1493
-
1494
- # -- Workflow routes ---
1495
-
1496
- @router.get(
1497
- "/workflows",
1498
- response_model=List[WorkflowSummaryResponse],
1499
- response_model_exclude_none=True,
1500
- tags=["Workflows"],
1501
- operation_id="get_workflows",
1502
- summary="List All Workflows",
1503
- description=(
1504
- "Retrieve a comprehensive list of all workflows configured in this OS instance.\n\n"
1505
- "**Return Information:**\n"
1506
- "- Workflow metadata (ID, name, description)\n"
1507
- "- Input schema requirements\n"
1508
- "- Step sequence and execution flow\n"
1509
- "- Associated agents and teams"
1510
- ),
1511
- responses={
1512
- 200: {
1513
- "description": "List of workflows retrieved successfully",
1514
- "content": {
1515
- "application/json": {
1516
- "example": [
1517
- {
1518
- "id": "content-creation-workflow",
1519
- "name": "Content Creation Workflow",
1520
- "description": "Automated content creation from blog posts to social media",
1521
- "db_id": "123",
1522
- }
1523
- ]
1524
- }
1525
- },
1526
- }
1527
- },
1528
- )
1529
- async def get_workflows() -> List[WorkflowSummaryResponse]:
1530
- if os.workflows is None:
1531
- return []
1532
-
1533
- return [WorkflowSummaryResponse.from_workflow(workflow) for workflow in os.workflows]
1534
-
1535
- @router.get(
1536
- "/workflows/{workflow_id}",
1537
- response_model=WorkflowResponse,
1538
- response_model_exclude_none=True,
1539
- tags=["Workflows"],
1540
- operation_id="get_workflow",
1541
- summary="Get Workflow Details",
1542
- description=("Retrieve detailed configuration and step information for a specific workflow."),
1543
- responses={
1544
- 200: {
1545
- "description": "Workflow details retrieved successfully",
1546
- "content": {
1547
- "application/json": {
1548
- "example": {
1549
- "id": "content-creation-workflow",
1550
- "name": "Content Creation Workflow",
1551
- "description": "Automated content creation from blog posts to social media",
1552
- "db_id": "123",
1553
- }
1554
- }
1555
- },
1556
- },
1557
- 404: {"description": "Workflow not found", "model": NotFoundResponse},
1558
- },
1559
- )
1560
- async def get_workflow(workflow_id: str) -> WorkflowResponse:
1561
- workflow = get_workflow_by_id(workflow_id, os.workflows)
1562
- if workflow is None:
1563
- raise HTTPException(status_code=404, detail="Workflow not found")
1564
-
1565
- return WorkflowResponse.from_workflow(workflow)
1566
-
1567
- @router.post(
1568
- "/workflows/{workflow_id}/runs",
1569
- tags=["Workflows"],
1570
- operation_id="create_workflow_run",
1571
- response_model_exclude_none=True,
1572
- summary="Execute Workflow",
1573
- description=(
1574
- "Execute a workflow with the provided input data. Workflows can run in streaming or batch mode.\n\n"
1575
- "**Execution Modes:**\n"
1576
- "- **Streaming (`stream=true`)**: Real-time step-by-step execution updates via SSE\n"
1577
- "- **Non-Streaming (`stream=false`)**: Complete workflow execution with final result\n\n"
1578
- "**Workflow Execution Process:**\n"
1579
- "1. Input validation against workflow schema\n"
1580
- "2. Sequential or parallel step execution based on workflow design\n"
1581
- "3. Data flow between steps with transformation\n"
1582
- "4. Error handling and automatic retries where configured\n"
1583
- "5. Final result compilation and response\n\n"
1584
- "**Session Management:**\n"
1585
- "Workflows support session continuity for stateful execution across multiple runs."
1586
- ),
1587
- responses={
1588
- 200: {
1589
- "description": "Workflow executed successfully",
1590
- "content": {
1591
- "text/event-stream": {
1592
- "example": 'event: RunStarted\ndata: {"content": "Hello!", "run_id": "123..."}\n\n'
1593
- },
1594
- },
1595
- },
1596
- 400: {"description": "Invalid input data or workflow configuration", "model": BadRequestResponse},
1597
- 404: {"description": "Workflow not found", "model": NotFoundResponse},
1598
- 500: {"description": "Workflow execution error", "model": InternalServerErrorResponse},
1599
- },
1600
- )
1601
- async def create_workflow_run(
1602
- workflow_id: str,
1603
- request: Request,
1604
- message: str = Form(...),
1605
- stream: bool = Form(True),
1606
- session_id: Optional[str] = Form(None),
1607
- user_id: Optional[str] = Form(None),
1608
- ):
1609
- kwargs = await _get_request_kwargs(request, create_workflow_run)
1610
-
1611
- if hasattr(request.state, "user_id"):
1612
- if user_id:
1613
- log_warning("User ID parameter passed in both request state and kwargs, using request state")
1614
- user_id = request.state.user_id
1615
- if hasattr(request.state, "session_id"):
1616
- if session_id:
1617
- log_warning("Session ID parameter passed in both request state and kwargs, using request state")
1618
- session_id = request.state.session_id
1619
- if hasattr(request.state, "session_state"):
1620
- session_state = request.state.session_state
1621
- if "session_state" in kwargs:
1622
- log_warning("Session state parameter passed in both request state and kwargs, using request state")
1623
- kwargs["session_state"] = session_state
1624
- if hasattr(request.state, "dependencies"):
1625
- dependencies = request.state.dependencies
1626
- if "dependencies" in kwargs:
1627
- log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
1628
- kwargs["dependencies"] = dependencies
1629
-
1630
- # Retrieve the workflow by ID
1631
- workflow = get_workflow_by_id(workflow_id, os.workflows)
1632
- if workflow is None:
1633
- raise HTTPException(status_code=404, detail="Workflow not found")
1634
-
1635
- if session_id:
1636
- logger.debug(f"Continuing session: {session_id}")
1637
- else:
1638
- logger.debug("Creating new session")
1639
- session_id = str(uuid4())
1640
-
1641
- # Return based on stream parameter
1642
- try:
1643
- if stream:
1644
- return StreamingResponse(
1645
- workflow_response_streamer(
1646
- workflow,
1647
- input=message,
1648
- session_id=session_id,
1649
- user_id=user_id,
1650
- **kwargs,
1651
- ),
1652
- media_type="text/event-stream",
1653
- )
233
+ async def migrate_database(db_id: str, target_version: Optional[str] = None):
234
+ db = await get_db(os.dbs, db_id)
235
+ if not db:
236
+ raise HTTPException(status_code=404, detail="Database not found")
237
+
238
+ if target_version:
239
+ # Use the session table as proxy for the database schema version
240
+ if isinstance(db, AsyncBaseDb):
241
+ current_version = await db.get_latest_schema_version(db.session_table_name)
1654
242
  else:
1655
- run_response = await workflow.arun(
1656
- input=message,
1657
- session_id=session_id,
1658
- user_id=user_id,
1659
- stream=False,
1660
- **kwargs,
1661
- )
1662
- return run_response.to_dict()
243
+ current_version = db.get_latest_schema_version(db.session_table_name)
1663
244
 
1664
- except InputCheckError as e:
1665
- raise HTTPException(status_code=400, detail=str(e))
1666
- except Exception as e:
1667
- # Handle unexpected runtime errors
1668
- raise HTTPException(status_code=500, detail=f"Error running workflow: {str(e)}")
1669
-
1670
- @router.post(
1671
- "/workflows/{workflow_id}/runs/{run_id}/cancel",
1672
- tags=["Workflows"],
1673
- operation_id="cancel_workflow_run",
1674
- summary="Cancel Workflow Run",
1675
- description=(
1676
- "Cancel a currently executing workflow run, stopping all active steps and cleanup.\n"
1677
- "**Note:** Complex workflows with multiple parallel steps may take time to fully cancel."
1678
- ),
1679
- responses={
1680
- 200: {},
1681
- 404: {"description": "Workflow or run not found", "model": NotFoundResponse},
1682
- 500: {"description": "Failed to cancel workflow run", "model": InternalServerErrorResponse},
1683
- },
1684
- )
1685
- async def cancel_workflow_run(workflow_id: str, run_id: str):
1686
- workflow = get_workflow_by_id(workflow_id, os.workflows)
1687
-
1688
- if workflow is None:
1689
- raise HTTPException(status_code=404, detail="Workflow not found")
245
+ if version.parse(target_version) > version.parse(current_version): # type: ignore
246
+ MigrationManager(db).up(target_version) # type: ignore
247
+ else:
248
+ MigrationManager(db).down(target_version) # type: ignore
1690
249
 
1691
- if not workflow.cancel_run(run_id=run_id):
1692
- raise HTTPException(status_code=500, detail="Failed to cancel run")
250
+ # If the target version is not provided, migrate to the latest version
251
+ else:
252
+ MigrationManager(db).up() # type: ignore
1693
253
 
1694
- return JSONResponse(content={}, status_code=200)
254
+ return JSONResponse(
255
+ content={"message": f"Database migrated successfully to version {target_version}"}, status_code=200
256
+ )
1695
257
 
1696
258
  return router