idun-agent-engine 0.2.7__py3-none-any.whl → 0.3.0__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 (41) hide show
  1. idun_agent_engine/_version.py +1 -1
  2. idun_agent_engine/agent/adk/__init__.py +5 -0
  3. idun_agent_engine/agent/adk/adk.py +296 -0
  4. idun_agent_engine/agent/base.py +7 -1
  5. idun_agent_engine/agent/haystack/haystack.py +5 -1
  6. idun_agent_engine/agent/langgraph/langgraph.py +146 -55
  7. idun_agent_engine/core/app_factory.py +9 -0
  8. idun_agent_engine/core/config_builder.py +214 -23
  9. idun_agent_engine/core/engine_config.py +1 -2
  10. idun_agent_engine/core/server_runner.py +2 -3
  11. idun_agent_engine/guardrails/__init__.py +0 -0
  12. idun_agent_engine/guardrails/base.py +24 -0
  13. idun_agent_engine/guardrails/guardrails_hub/guardrails_hub.py +101 -0
  14. idun_agent_engine/guardrails/guardrails_hub/utils.py +1 -0
  15. idun_agent_engine/mcp/__init__.py +5 -0
  16. idun_agent_engine/mcp/helpers.py +97 -0
  17. idun_agent_engine/mcp/registry.py +109 -0
  18. idun_agent_engine/observability/__init__.py +6 -2
  19. idun_agent_engine/observability/base.py +73 -12
  20. idun_agent_engine/observability/gcp_logging/__init__.py +0 -0
  21. idun_agent_engine/observability/gcp_logging/gcp_logging_handler.py +52 -0
  22. idun_agent_engine/observability/gcp_trace/__init__.py +0 -0
  23. idun_agent_engine/observability/gcp_trace/gcp_trace_handler.py +116 -0
  24. idun_agent_engine/observability/langfuse/langfuse_handler.py +17 -10
  25. idun_agent_engine/server/dependencies.py +13 -1
  26. idun_agent_engine/server/lifespan.py +83 -16
  27. idun_agent_engine/server/routers/agent.py +116 -24
  28. idun_agent_engine/server/routers/agui.py +47 -0
  29. idun_agent_engine/server/routers/base.py +55 -1
  30. idun_agent_engine/templates/__init__.py +1 -0
  31. idun_agent_engine/templates/correction.py +65 -0
  32. idun_agent_engine/templates/deep_research.py +40 -0
  33. idun_agent_engine/templates/translation.py +70 -0
  34. {idun_agent_engine-0.2.7.dist-info → idun_agent_engine-0.3.0.dist-info}/METADATA +62 -10
  35. idun_agent_engine-0.3.0.dist-info/RECORD +60 -0
  36. {idun_agent_engine-0.2.7.dist-info → idun_agent_engine-0.3.0.dist-info}/WHEEL +1 -1
  37. idun_platform_cli/groups/agent/package.py +3 -3
  38. idun_platform_cli/groups/agent/serve.py +8 -5
  39. idun_agent_engine/cli/__init__.py +0 -16
  40. idun_agent_engine-0.2.7.dist-info/RECORD +0 -43
  41. {idun_agent_engine-0.2.7.dist-info → idun_agent_engine-0.3.0.dist-info}/entry_points.txt +0 -0
@@ -1,3 +1,3 @@
1
1
  """Version information for Idun Agent Engine."""
2
2
 
3
- __version__ = "0.2.7"
3
+ __version__ = "0.3.0"
@@ -0,0 +1,5 @@
1
+ """ADK Agent implementation."""
2
+
3
+ from .adk import AdkAgent
4
+
5
+ __all__ = ["AdkAgent"]
@@ -0,0 +1,296 @@
1
+ """ADK agent adapter implementing the BaseAgent protocol."""
2
+
3
+ import importlib.util
4
+ import uuid
5
+ from collections.abc import AsyncGenerator
6
+ from typing import Any
7
+
8
+ from google.adk.apps.app import App
9
+ from google.adk.memory import (
10
+ InMemoryMemoryService,
11
+ VertexAiMemoryBankService,
12
+ )
13
+ from google.adk.sessions import (
14
+ DatabaseSessionService,
15
+ InMemorySessionService,
16
+ VertexAiSessionService,
17
+ )
18
+ from idun_agent_schema.engine.adk import (
19
+ AdkAgentConfig,
20
+ AdkDatabaseSessionConfig,
21
+ AdkInMemoryMemoryConfig,
22
+ AdkInMemorySessionConfig,
23
+ AdkVertexAiMemoryConfig,
24
+ AdkVertexAiSessionConfig,
25
+ )
26
+ from idun_agent_schema.engine.observability_v2 import ObservabilityConfig
27
+
28
+ from ag_ui_adk import ADKAgent as ADKAGUIAgent
29
+ from idun_agent_engine.agent import base as agent_base
30
+ from idun_agent_engine import observability
31
+
32
+
33
+ class AdkAgent(agent_base.BaseAgent):
34
+ """ADK agent adapter implementing the BaseAgent protocol."""
35
+
36
+ def __init__(self):
37
+ """Initialize an unconfigured AdkAgent with default state."""
38
+ self._id = str(uuid.uuid4())
39
+ self._agent_type = "ADK"
40
+ self._agent_instance: Any = None
41
+ self._copilotkit_agent_instance: ADKAGUIAgent | None = None
42
+ self._configuration: AdkAgentConfig | None = None
43
+ self._name: str = "Unnamed ADK Agent"
44
+ self._infos: dict[str, Any] = {
45
+ "status": "Uninitialized",
46
+ "name": self._name,
47
+ "id": self._id,
48
+ }
49
+ self._session_service: Any = None
50
+ self._memory_service: Any = None
51
+ # Observability (provider-agnostic)
52
+ self._obs_callbacks: list[Any] | None = None
53
+
54
+ @property
55
+ def id(self) -> str:
56
+ """Return unique identifier for this agent instance."""
57
+ return self._id
58
+
59
+ @property
60
+ def agent_type(self) -> str:
61
+ """Return agent type label."""
62
+ return self._agent_type
63
+
64
+ @property
65
+ def name(self) -> str:
66
+ """Return configured human-readable agent name."""
67
+ return self._name
68
+
69
+ @property
70
+ def agent_instance(self) -> Any:
71
+ """Return the underlying ADK agent instance.
72
+
73
+ Raises:
74
+ RuntimeError: If the agent is not yet initialized.
75
+ """
76
+ if self._agent_instance is None:
77
+ raise RuntimeError("Agent not initialized. Call initialize() first.")
78
+ return self._agent_instance
79
+
80
+ @property
81
+ def copilotkit_agent_instance(self) -> ADKAGUIAgent:
82
+ """Return the CopilotKit agent instance.
83
+
84
+ Raises:
85
+ RuntimeError: If the CopilotKit agent is not yet initialized.
86
+ """
87
+ if self._copilotkit_agent_instance is None:
88
+ raise RuntimeError(
89
+ "CopilotKit agent not initialized. Call initialize() first."
90
+ )
91
+ return self._copilotkit_agent_instance
92
+
93
+ @property
94
+ def configuration(self) -> AdkAgentConfig:
95
+ """Return validated configuration.
96
+
97
+ Raises:
98
+ RuntimeError: If the agent has not been configured yet.
99
+ """
100
+ if not self._configuration:
101
+ raise RuntimeError("Agent not configured. Call initialize() first.")
102
+ return self._configuration
103
+
104
+ @property
105
+ def infos(self) -> dict[str, Any]:
106
+ """Return diagnostic information about the agent instance."""
107
+ self._infos["underlying_agent_type"] = (
108
+ str(type(self._agent_instance)) if self._agent_instance else "N/A"
109
+ )
110
+ return self._infos
111
+
112
+ async def initialize(
113
+ self,
114
+ config: AdkAgentConfig,
115
+ observability_config: list[ObservabilityConfig] | None = None,
116
+ ) -> None:
117
+ """Initialize the ADK agent asynchronously."""
118
+ self._configuration = AdkAgentConfig.model_validate(config)
119
+
120
+ self._name = self._configuration.app_name or "Unnamed ADK Agent"
121
+ self._infos["name"] = self._name
122
+
123
+ # Observability (provider-agnostic)
124
+ if observability_config:
125
+ handlers, infos = observability.create_observability_handlers(
126
+ observability_config # type: ignore[arg-type]
127
+ )
128
+ self._obs_callbacks = []
129
+ for handler in handlers:
130
+ # Even if callbacks aren't used by ADK directly, instantiating the handler
131
+ # might set up global instrumentation (e.g. Phoenix, Langfuse env vars).
132
+ self._obs_callbacks.extend(handler.get_callbacks())
133
+
134
+ if infos:
135
+ self._infos["observability"] = infos
136
+
137
+ if observability_config:
138
+ try:
139
+ # Check if langfuse is enabled in any of the observability configs
140
+ def _is_langfuse_provider(c: Any) -> bool:
141
+ provider = getattr(c, "provider", None)
142
+ if provider is None and isinstance(c, dict):
143
+ provider = c.get("provider")
144
+
145
+ if provider is not None and hasattr(provider, "value"):
146
+ provider = provider.value
147
+
148
+ return str(provider).lower() == "langfuse"
149
+
150
+ is_langfuse_enabled = any(
151
+ _is_langfuse_provider(config) for config in observability_config
152
+ )
153
+
154
+ if is_langfuse_enabled:
155
+ import os
156
+ langfuse_pk = os.environ.get("LANGFUSE_PUBLIC_KEY")
157
+ langfuse_host = os.environ.get("LANGFUSE_BASE_URL")
158
+ print(f"LANGFUSE_PUBLIC_KEY: {langfuse_pk}")
159
+ print(f"LANGFUSE_BASE_URL: {langfuse_host}")
160
+ try:
161
+ from openinference.instrumentation.google_adk import (
162
+ GoogleADKInstrumentor,
163
+ )
164
+
165
+ GoogleADKInstrumentor().instrument()
166
+ print("GoogleADKInstrumentor instrumented successfully.")
167
+ except ImportError:
168
+ print(
169
+ "openinference-instrumentation-google-adk not installed, skipping Google ADK instrumentation."
170
+ )
171
+ except Exception as e:
172
+ print(f"Failed to instrument Google ADK: {e}")
173
+ except Exception as e:
174
+ print(f"Error checking observability config for ADK instrumentation: {e}")
175
+
176
+ # Initialize Session Service
177
+ await self._initialize_session_service()
178
+
179
+ # Initialize Memory Service
180
+ await self._initialize_memory_service()
181
+
182
+ # Load the agent instance
183
+ agent = self._load_agent(self._configuration.agent)
184
+
185
+ self._agent_instance = App(root_agent=agent, name=self._name)
186
+
187
+ # Initialize CopilotKit/AG-UI Agent Wrapper
188
+ # TODO: Pass session and memory services when supported by AG-UI ADK adapter if needed
189
+ self._copilotkit_agent_instance = ADKAGUIAgent(
190
+ adk_agent=agent,
191
+ session_service=self._session_service,
192
+ memory_service=self._memory_service,
193
+ app_name=self._name,
194
+ )
195
+
196
+ self._infos["status"] = "Initialized"
197
+ self._infos["config_used"] = self._configuration.model_dump()
198
+
199
+ async def _initialize_session_service(self) -> None:
200
+ """Initialize the session service based on configuration."""
201
+ if not self._configuration:
202
+ raise RuntimeError("Configuration not initialized")
203
+
204
+ if not self._configuration.session_service:
205
+ # Default to InMemory if not specified
206
+ self._session_service = InMemorySessionService()
207
+ return
208
+
209
+ config = self._configuration.session_service
210
+ if isinstance(config, AdkInMemorySessionConfig):
211
+ self._session_service = InMemorySessionService()
212
+ elif isinstance(config, AdkVertexAiSessionConfig):
213
+ self._session_service = VertexAiSessionService(
214
+ project=config.project_id,
215
+ location=config.location,
216
+ agent_engine_id=config.reasoning_engine_app_name,
217
+ )
218
+ elif isinstance(config, AdkDatabaseSessionConfig):
219
+ self._session_service = DatabaseSessionService(db_url=config.db_url)
220
+ else:
221
+ raise ValueError(f"Unsupported session service type: {config.type}") # type: ignore
222
+
223
+ async def _initialize_memory_service(self) -> None:
224
+ """Initialize the memory service based on configuration."""
225
+ if not self._configuration:
226
+ raise RuntimeError("Configuration not initialized")
227
+
228
+ if not self._configuration.memory_service:
229
+ # Default to InMemory if not specified
230
+ self._memory_service = InMemoryMemoryService()
231
+ return
232
+
233
+ config = self._configuration.memory_service
234
+ if isinstance(config, AdkInMemoryMemoryConfig):
235
+ self._memory_service = InMemoryMemoryService()
236
+ elif isinstance(config, AdkVertexAiMemoryConfig):
237
+ self._memory_service = VertexAiMemoryBankService(
238
+ project=config.project_id,
239
+ location=config.location,
240
+ agent_engine_id=config.memory_bank_id,
241
+ )
242
+ else:
243
+ raise ValueError(f"Unsupported memory service type: {config.type}") # type: ignore
244
+
245
+ def _load_agent(self, agent_definition: str) -> Any:
246
+ """Loads an agent instance from a specified path."""
247
+ try:
248
+ module_path, agent_variable_name = agent_definition.rsplit(":", 1)
249
+ except ValueError:
250
+ raise ValueError(
251
+ "agent_definition must be in the format 'path/to/file.py:variable_name'"
252
+ ) from None
253
+
254
+ try:
255
+ from pathlib import Path
256
+
257
+ resolved_path = Path(module_path).resolve()
258
+ spec = importlib.util.spec_from_file_location(
259
+ agent_variable_name, str(resolved_path)
260
+ )
261
+ if spec is None or spec.loader is None:
262
+ raise ImportError(f"Could not load spec for module at {module_path}")
263
+
264
+ module = importlib.util.module_from_spec(spec)
265
+ spec.loader.exec_module(module)
266
+
267
+ agent_instance = getattr(module, agent_variable_name)
268
+ return agent_instance
269
+ except (FileNotFoundError, ImportError, AttributeError) as e:
270
+ raise ValueError(
271
+ f"Failed to load agent from {agent_definition}: {e}"
272
+ ) from e
273
+
274
+ async def invoke(self, message: Any) -> Any:
275
+ """Process a single input to chat with the agent."""
276
+ if self._agent_instance is None:
277
+ raise RuntimeError(
278
+ "Agent not initialized. Call initialize() before processing messages."
279
+ )
280
+
281
+ # TODO: Implement ADK invoke logic using session and memory services
282
+ raise NotImplementedError("ADK invoke not implemented yet")
283
+
284
+ async def stream(self, message: Any) -> AsyncGenerator[Any]:
285
+ """Process a single input message and return an asynchronous stream."""
286
+ if self._agent_instance is None:
287
+ raise RuntimeError(
288
+ "Agent not initialized. Call initialize() before processing messages."
289
+ )
290
+
291
+ # TODO: Implement ADK stream logic using session and memory services
292
+ raise NotImplementedError("ADK stream not implemented yet")
293
+
294
+ # Required to make this a generator
295
+ if False:
296
+ yield
@@ -8,6 +8,7 @@ from collections.abc import AsyncGenerator
8
8
  from typing import Any
9
9
 
10
10
  from idun_agent_schema.engine.agent import BaseAgentConfig
11
+ from idun_agent_schema.engine.observability_v2 import ObservabilityConfig
11
12
 
12
13
 
13
14
  class BaseAgent[ConfigType: BaseAgentConfig](ABC):
@@ -63,13 +64,18 @@ class BaseAgent[ConfigType: BaseAgentConfig](ABC):
63
64
  pass
64
65
 
65
66
  @abstractmethod
66
- async def initialize(self, config: dict[str, Any]) -> None:
67
+ async def initialize(
68
+ self,
69
+ config: dict[str, Any],
70
+ observability: list[ObservabilityConfig] | None = None,
71
+ ) -> None:
67
72
  """Initialize the agent with a given configuration.
68
73
 
69
74
  This method should set up the underlying agent framework instance.
70
75
 
71
76
  Args:
72
77
  config: A dictionary containing the agent's configuration.
78
+ observability: Optional list of observability configurations.
73
79
  """
74
80
  pass
75
81
 
@@ -126,7 +126,11 @@ class HaystackAgent(BaseAgent):
126
126
  self._langfuse_tracing = True
127
127
  logger.info("Agent tracing not supported yet")
128
128
 
129
- async def initialize(self, config: HaystackAgentConfig | dict[str, Any]) -> None:
129
+ async def initialize(
130
+ self,
131
+ config: HaystackAgentConfig | dict[str, Any],
132
+ observability_config: list[ObservabilityConfig] | None = None,
133
+ ) -> None:
130
134
  try:
131
135
  logger.debug(f"Initializing haystack agent config: {config}...")
132
136
 
@@ -1,6 +1,7 @@
1
1
  """LangGraph agent adapter implementing the BaseAgent protocol."""
2
2
 
3
3
  import importlib.util
4
+ import importlib
4
5
  import uuid
5
6
  from collections.abc import AsyncGenerator
6
7
  from typing import Any
@@ -9,11 +10,17 @@ import aiosqlite
9
10
  from ag_ui.core import events as ag_events
10
11
  from ag_ui.core import types as ag_types
11
12
  from idun_agent_schema.engine.langgraph import (
13
+ InMemoryCheckpointConfig,
12
14
  LangGraphAgentConfig,
15
+ PostgresCheckpointConfig,
13
16
  SqliteCheckpointConfig,
14
17
  )
18
+ from idun_agent_schema.engine.observability_v2 import ObservabilityConfig
19
+ from langgraph.checkpoint.memory import InMemorySaver
20
+ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
15
21
  from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
16
22
  from langgraph.graph import StateGraph
23
+ from langgraph.graph.state import CompiledStateGraph
17
24
 
18
25
  from idun_agent_engine import observability
19
26
  from idun_agent_engine.agent import base as agent_base
@@ -111,7 +118,11 @@ class LanggraphAgent(agent_base.BaseAgent):
111
118
  )
112
119
  return self._infos
113
120
 
114
- async def initialize(self, config: LangGraphAgentConfig) -> None:
121
+ async def initialize(
122
+ self,
123
+ config: LangGraphAgentConfig,
124
+ observability_config: list[ObservabilityConfig] | None = None,
125
+ ) -> None:
115
126
  """Initialize the LangGraph agent asynchronously."""
116
127
  self._configuration = LangGraphAgentConfig.model_validate(config)
117
128
 
@@ -120,55 +131,83 @@ class LanggraphAgent(agent_base.BaseAgent):
120
131
 
121
132
  await self._setup_persistence()
122
133
 
123
- # Observability (provider-agnostic). Prefer generic block; fallback to legacy langfuse block.
124
- obs_cfg = None
125
- try:
126
- if getattr(self._configuration, "observability", None):
127
- obs_cfg = self._configuration.observability.resolved() # type: ignore[attr-defined]
128
- elif getattr(self._configuration, "langfuse", None):
129
- lf = self._configuration.langfuse.resolved() # type: ignore[attr-defined]
130
- obs_cfg = type(
131
- "_Temp",
132
- (),
133
- {
134
- "provider": "langfuse",
135
- "enabled": lf.enabled,
136
- "options": {
137
- "host": lf.host,
138
- "public_key": lf.public_key,
139
- "secret_key": lf.secret_key,
140
- "run_name": lf.run_name,
141
- },
142
- },
143
- )()
144
- except Exception:
145
- obs_cfg = None
146
-
147
- if obs_cfg and getattr(obs_cfg, "enabled", False):
148
- provider = getattr(obs_cfg, "provider", None)
149
- options = dict(getattr(obs_cfg, "options", {}) or {})
150
- # Fallback: if using Langfuse and run_name is not provided, use agent name
151
- if provider == "langfuse" and not options.get("run_name"):
152
- options["run_name"] = self._name
153
-
154
- handler, info = observability.create_observability_handler(
155
- {
156
- "provider": provider,
157
- "enabled": True,
158
- "options": options,
159
- }
134
+ # Observability (provider-agnostic)
135
+ if observability_config:
136
+ handlers, infos = observability.create_observability_handlers(
137
+ observability_config # type: ignore[arg-type]
160
138
  )
161
- if handler:
162
- self._obs_callbacks = handler.get_callbacks()
163
- self._obs_run_name = handler.get_run_name()
164
- if info:
165
- self._infos["observability"] = dict(info)
139
+ self._obs_callbacks = []
140
+ for handler in handlers:
141
+ self._obs_callbacks.extend(handler.get_callbacks())
142
+ # Use the first run name found if not set
143
+ if not self._obs_run_name:
144
+ self._obs_run_name = handler.get_run_name()
145
+
146
+ if infos:
147
+ self._infos["observability"] = infos
148
+
149
+ # Fallback to legacy generic block or langfuse block if no new observability config provided
150
+ elif getattr(self._configuration, "observability", None) or getattr(
151
+ self._configuration, "langfuse", None
152
+ ):
153
+ obs_cfg = None
154
+ try:
155
+ if getattr(self._configuration, "observability", None):
156
+ obs_cfg = self._configuration.observability.resolved() # type: ignore[attr-defined]
157
+ elif getattr(self._configuration, "langfuse", None):
158
+ lf = self._configuration.langfuse.resolved() # type: ignore[attr-defined]
159
+ obs_cfg = type(
160
+ "_Temp",
161
+ (),
162
+ {
163
+ "provider": "langfuse",
164
+ "enabled": lf.enabled,
165
+ "options": {
166
+ "host": lf.host,
167
+ "public_key": lf.public_key,
168
+ "secret_key": lf.secret_key,
169
+ "run_name": lf.run_name,
170
+ },
171
+ },
172
+ )()
173
+ except Exception:
174
+ obs_cfg = None
175
+
176
+ if obs_cfg and getattr(obs_cfg, "enabled", False):
177
+ provider = getattr(obs_cfg, "provider", None)
178
+ options = dict(getattr(obs_cfg, "options", {}) or {})
179
+ # Fallback: if using Langfuse and run_name is not provided, use agent name
180
+ if provider == "langfuse" and not options.get("run_name"):
181
+ options["run_name"] = self._name
182
+
183
+ handler, info = observability.create_observability_handler(
184
+ {
185
+ "provider": provider,
186
+ "enabled": True,
187
+ "options": options,
188
+ }
189
+ )
190
+ if handler:
191
+ self._obs_callbacks = handler.get_callbacks()
192
+ self._obs_run_name = handler.get_run_name()
193
+ if info:
194
+ self._infos["observability"] = dict(info)
166
195
 
167
196
  graph_builder = self._load_graph_builder(self._configuration.graph_definition)
168
197
  self._infos["graph_definition"] = self._configuration.graph_definition
169
198
 
170
- self._agent_instance = graph_builder.compile(
171
- checkpointer=self._checkpointer, store=self._store
199
+ if isinstance(graph_builder, StateGraph):
200
+ self._agent_instance = graph_builder.compile(
201
+ checkpointer=self._checkpointer, store=self._store
202
+ )
203
+ elif isinstance(graph_builder, CompiledStateGraph):
204
+ self._agent_instance = graph_builder
205
+
206
+ self._copilotkit_agent_instance = LangGraphAGUIAgent(
207
+ name=self._name,
208
+ description="Agent description", # TODO: add agent description
209
+ graph=self._agent_instance,
210
+ config={"callbacks": self._obs_callbacks} if self._obs_callbacks else None,
172
211
  )
173
212
 
174
213
  self._copilotkit_agent_instance = LangGraphAGUIAgent(
@@ -178,10 +217,18 @@ class LanggraphAgent(agent_base.BaseAgent):
178
217
  )
179
218
 
180
219
  if self._agent_instance:
181
- self._input_schema = self._agent_instance.input_schema
182
- self._output_schema = self._agent_instance.output_schema
183
- self._infos["input_schema"] = str(self._input_schema)
184
- self._infos["output_schema"] = str(self._output_schema)
220
+ try:
221
+ self._input_schema = self._agent_instance.input_schema
222
+ self._output_schema = self._agent_instance.output_schema
223
+ self._infos["input_schema"] = str(self._input_schema)
224
+ self._infos["output_schema"] = str(self._output_schema)
225
+ except Exception:
226
+ print("Could not parse schema")
227
+ self._input_schema = self._configuration.input_schema_definition
228
+ self._output_schema = self._configuration.output_schema_definition
229
+ self._infos["input_schema"] = "Cannot extract schema"
230
+ self._infos["output_schema"] = "Cannot extract schema"
231
+
185
232
  else:
186
233
  self._input_schema = self._configuration.input_schema_definition
187
234
  self._output_schema = self._configuration.output_schema_definition
@@ -210,8 +257,23 @@ class LanggraphAgent(agent_base.BaseAgent):
210
257
  self._infos["checkpointer"] = (
211
258
  self._configuration.checkpointer.model_dump()
212
259
  )
260
+ elif isinstance(self._configuration.checkpointer, InMemoryCheckpointConfig):
261
+ self._checkpointer = InMemorySaver()
262
+ self._infos["checkpointer"] = (
263
+ self._configuration.checkpointer.model_dump()
264
+ )
265
+ elif isinstance(self._configuration.checkpointer, PostgresCheckpointConfig):
266
+ self._checkpointer = AsyncPostgresSaver.from_conn_string(
267
+ self._configuration.checkpointer.db_url
268
+ )
269
+ await self._checkpointer.setup()
270
+ self._infos["checkpointer"] = (
271
+ self._configuration.checkpointer.model_dump()
272
+ )
213
273
  else:
214
- raise NotImplementedError("Only SQLite checkpointer is supported.")
274
+ raise NotImplementedError(
275
+ f"Checkpointer type {type(self._configuration.checkpointer)} is not supported."
276
+ )
215
277
 
216
278
  if self._configuration.store:
217
279
  raise NotImplementedError("Store functionality is not yet implemented.")
@@ -220,14 +282,24 @@ class LanggraphAgent(agent_base.BaseAgent):
220
282
  """Loads a StateGraph instance from a specified path."""
221
283
  try:
222
284
  module_path, graph_variable_name = graph_definition.rsplit(":", 1)
285
+ if not module_path.endswith(".py"):
286
+ module_path += ".py"
223
287
  except ValueError:
224
288
  raise ValueError(
225
289
  "graph_definition must be in the format 'path/to/file.py:variable_name'"
226
290
  ) from None
227
291
 
292
+ # Try loading as a file path first
228
293
  try:
294
+ import os
295
+ print("Current directory: ", os.getcwd()) # TODO remove
229
296
  from pathlib import Path
297
+
230
298
  resolved_path = Path(module_path).resolve()
299
+ # If the file doesn't exist, it might be a python module path
300
+ if not resolved_path.exists():
301
+ raise FileNotFoundError
302
+
231
303
  spec = importlib.util.spec_from_file_location(
232
304
  graph_variable_name, str(resolved_path)
233
305
  )
@@ -238,17 +310,36 @@ class LanggraphAgent(agent_base.BaseAgent):
238
310
  spec.loader.exec_module(module)
239
311
 
240
312
  graph_builder = getattr(module, graph_variable_name)
241
- except (FileNotFoundError, ImportError, AttributeError) as e:
242
- raise ValueError(
313
+ return self._validate_graph_builder(graph_builder, module_path, graph_variable_name)
314
+
315
+ except (FileNotFoundError, ImportError):
316
+ # Fallback: try loading as a python module
317
+ try:
318
+ module = importlib.import_module(module_path)
319
+ graph_builder = getattr(module, graph_variable_name)
320
+ return self._validate_graph_builder(graph_builder, module_path, graph_variable_name)
321
+ except ImportError as e:
322
+ raise ValueError(
323
+ f"Failed to load agent from {graph_definition}. Checked file path and python module: {e}"
324
+ ) from e
325
+ except AttributeError as e:
326
+ raise ValueError(
327
+ f"Variable '{graph_variable_name}' not found in module {module_path}: {e}"
328
+ ) from e
329
+ except Exception as e:
330
+ raise ValueError(
243
331
  f"Failed to load agent from {graph_definition}: {e}"
244
332
  ) from e
245
333
 
246
- if not isinstance(graph_builder, StateGraph):
334
+ def _validate_graph_builder(self, graph_builder: Any, module_path: str, graph_variable_name: str) -> StateGraph:
335
+ # TODO to remove, dirty fix for template deepagent langgraph
336
+ if not isinstance(graph_builder, StateGraph) and not isinstance(
337
+ graph_builder, CompiledStateGraph
338
+ ):
247
339
  raise TypeError(
248
340
  f"The variable '{graph_variable_name}' from {module_path} is not a StateGraph instance."
249
341
  )
250
-
251
- return graph_builder
342
+ return graph_builder # type: ignore[return-value]
252
343
 
253
344
  async def invoke(self, message: Any) -> Any:
254
345
  """Process a single input to chat with the agent.
@@ -8,6 +8,7 @@ setting up routes, dependencies, and lifecycle management behind the scenes.
8
8
  from typing import Any
9
9
 
10
10
  from fastapi import FastAPI
11
+ from fastapi.middleware.cors import CORSMiddleware
11
12
 
12
13
  from ..server.lifespan import lifespan
13
14
  from ..server.routers.agent import agent_router
@@ -54,6 +55,14 @@ def create_app(
54
55
  redoc_url="/redoc",
55
56
  )
56
57
 
58
+ app.add_middleware(
59
+ CORSMiddleware,
60
+ allow_origins=["*"],
61
+ allow_credentials=True,
62
+ allow_methods=["*"],
63
+ allow_headers=["*"],
64
+ )
65
+
57
66
  # Store configuration in app state for lifespan to use
58
67
  app.state.engine_config = validated_config
59
68