idun-agent-engine 0.2.6__py3-none-any.whl → 0.3.0__py3-none-any.whl

idun_agent_engine/agent/langgraph/langgraph.py

@@ -1,6 +1,7 @@
 """LangGraph agent adapter implementing the BaseAgent protocol."""
 
 import importlib.util
+import importlib
 import uuid
 from collections.abc import AsyncGenerator
 from typing import Any
@@ -9,14 +10,21 @@ import aiosqlite
 from ag_ui.core import events as ag_events
 from ag_ui.core import types as ag_types
 from idun_agent_schema.engine.langgraph import (
+    InMemoryCheckpointConfig,
     LangGraphAgentConfig,
+    PostgresCheckpointConfig,
     SqliteCheckpointConfig,
 )
+from idun_agent_schema.engine.observability_v2 import ObservabilityConfig
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
 from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
 
 from idun_agent_engine import observability
 from idun_agent_engine.agent import base as agent_base
+from copilotkit import LangGraphAGUIAgent
 
 
 class LanggraphAgent(agent_base.BaseAgent):
@@ -29,6 +37,7 @@ class LanggraphAgent(agent_base.BaseAgent):
         self._input_schema: Any = None
         self._output_schema: Any = None
         self._agent_instance: Any = None
+        self._copilotkit_agent_instance: LangGraphAGUIAgent | None = None
         self._checkpointer: Any = None
         self._store: Any = None
         self._connection: Any = None
@@ -79,6 +88,17 @@ class LanggraphAgent(agent_base.BaseAgent):
             raise RuntimeError("Agent not initialized. Call initialize() first.")
         return self._agent_instance
 
+    @property
+    def copilotkit_agent_instance(self) -> LangGraphAGUIAgent:
+        """Return the CopilotKit agent instance.
+
+        Raises:
+            RuntimeError: If the CopilotKit agent is not yet initialized.
+        """
+        if self._copilotkit_agent_instance is None:
+            raise RuntimeError("CopilotKit agent not initialized. Call initialize() first.")
+        return self._copilotkit_agent_instance
+
     @property
     def configuration(self) -> LangGraphAgentConfig:
         """Return validated configuration.
@@ -98,7 +118,11 @@ class LanggraphAgent(agent_base.BaseAgent):
         )
         return self._infos
 
-    async def initialize(self, config: LangGraphAgentConfig) -> None:
+    async def initialize(
+        self,
+        config: LangGraphAgentConfig,
+        observability_config: list[ObservabilityConfig] | None = None,
+    ) -> None:
         """Initialize the LangGraph agent asynchronously."""
         self._configuration = LangGraphAgentConfig.model_validate(config)
 
@@ -107,62 +131,104 @@ class LanggraphAgent(agent_base.BaseAgent):
 
         await self._setup_persistence()
 
-        # Observability (provider-agnostic). Prefer generic block; fallback to legacy langfuse block.
-        obs_cfg = None
-        try:
-            if getattr(self._configuration, "observability", None):
-                obs_cfg = self._configuration.observability.resolved()  # type: ignore[attr-defined]
-            elif getattr(self._configuration, "langfuse", None):
-                lf = self._configuration.langfuse.resolved()  # type: ignore[attr-defined]
-                obs_cfg = type(
-                    "_Temp",
-                    (),
-                    {
-                        "provider": "langfuse",
-                        "enabled": lf.enabled,
-                        "options": {
-                            "host": lf.host,
-                            "public_key": lf.public_key,
-                            "secret_key": lf.secret_key,
-                            "run_name": lf.run_name,
-                        },
-                    },
-                )()
-        except Exception:
-            obs_cfg = None
-
-        if obs_cfg and getattr(obs_cfg, "enabled", False):
-            provider = getattr(obs_cfg, "provider", None)
-            options = dict(getattr(obs_cfg, "options", {}) or {})
-            # Fallback: if using Langfuse and run_name is not provided, use agent name
-            if provider == "langfuse" and not options.get("run_name"):
-                options["run_name"] = self._name
-
-            handler, info = observability.create_observability_handler(
-                {
-                    "provider": provider,
-                    "enabled": True,
-                    "options": options,
-                }
+        # Observability (provider-agnostic)
+        if observability_config:
+            handlers, infos = observability.create_observability_handlers(
+                observability_config # type: ignore[arg-type]
             )
-            if handler:
-                self._obs_callbacks = handler.get_callbacks()
-                self._obs_run_name = handler.get_run_name()
-            if info:
-                self._infos["observability"] = dict(info)
+            self._obs_callbacks = []
+            for handler in handlers:
+                self._obs_callbacks.extend(handler.get_callbacks())
+                # Use the first run name found if not set
+                if not self._obs_run_name:
+                    self._obs_run_name = handler.get_run_name()
+
+            if infos:
+                self._infos["observability"] = infos
+
+        # Fallback to legacy generic block or langfuse block if no new observability config provided
+        elif getattr(self._configuration, "observability", None) or getattr(
+            self._configuration, "langfuse", None
+        ):
+            obs_cfg = None
+            try:
+                if getattr(self._configuration, "observability", None):
+                    obs_cfg = self._configuration.observability.resolved()  # type: ignore[attr-defined]
+                elif getattr(self._configuration, "langfuse", None):
+                    lf = self._configuration.langfuse.resolved()  # type: ignore[attr-defined]
+                    obs_cfg = type(
+                        "_Temp",
+                        (),
+                        {
+                            "provider": "langfuse",
+                            "enabled": lf.enabled,
+                            "options": {
+                                "host": lf.host,
+                                "public_key": lf.public_key,
+                                "secret_key": lf.secret_key,
+                                "run_name": lf.run_name,
+                            },
+                        },
+                    )()
+            except Exception:
+                obs_cfg = None
+
+            if obs_cfg and getattr(obs_cfg, "enabled", False):
+                provider = getattr(obs_cfg, "provider", None)
+                options = dict(getattr(obs_cfg, "options", {}) or {})
+                # Fallback: if using Langfuse and run_name is not provided, use agent name
+                if provider == "langfuse" and not options.get("run_name"):
+                    options["run_name"] = self._name
+
+                handler, info = observability.create_observability_handler(
+                    {
+                        "provider": provider,
+                        "enabled": True,
+                        "options": options,
+                    }
+                )
+                if handler:
+                    self._obs_callbacks = handler.get_callbacks()
+                    self._obs_run_name = handler.get_run_name()
+                if info:
+                    self._infos["observability"] = dict(info)
 
         graph_builder = self._load_graph_builder(self._configuration.graph_definition)
         self._infos["graph_definition"] = self._configuration.graph_definition
 
-        self._agent_instance = graph_builder.compile(
-            checkpointer=self._checkpointer, store=self._store
+        if isinstance(graph_builder, StateGraph):
+            self._agent_instance = graph_builder.compile(
+                checkpointer=self._checkpointer, store=self._store
+            )
+        elif isinstance(graph_builder, CompiledStateGraph):
+            self._agent_instance = graph_builder
+
+        self._copilotkit_agent_instance = LangGraphAGUIAgent(
+            name=self._name,
+            description="Agent description", # TODO: add agent description
+            graph=self._agent_instance,
+            config={"callbacks": self._obs_callbacks} if self._obs_callbacks else None,
+        )
+
+        self._copilotkit_agent_instance = LangGraphAGUIAgent(
+            name=self._name,
+            description="Agent description", # TODO: add agent description
+            graph=self._agent_instance,
         )
 
         if self._agent_instance:
-            self._input_schema = self._agent_instance.input_schema
-            self._output_schema = self._agent_instance.output_schema
-            self._infos["input_schema"] = str(self._input_schema)
-            self._infos["output_schema"] = str(self._output_schema)
+            try:
+                self._input_schema = self._agent_instance.input_schema
+                self._output_schema = self._agent_instance.output_schema
+                self._infos["input_schema"] = str(self._input_schema)
+                self._infos["output_schema"] = str(self._output_schema)
+            except Exception:
+                print("Could not parse schema")
+                self._input_schema = self._configuration.input_schema_definition
+                self._output_schema = self._configuration.output_schema_definition
+                self._infos["input_schema"] = "Cannot extract schema"
+                self._infos["output_schema"] = "Cannot extract schema"
+
         else:
             self._input_schema = self._configuration.input_schema_definition
             self._output_schema = self._configuration.output_schema_definition
@@ -191,8 +257,23 @@ class LanggraphAgent(agent_base.BaseAgent):
                 self._infos["checkpointer"] = (
                     self._configuration.checkpointer.model_dump()
                 )
+            elif isinstance(self._configuration.checkpointer, InMemoryCheckpointConfig):
+                self._checkpointer = InMemorySaver()
+                self._infos["checkpointer"] = (
+                    self._configuration.checkpointer.model_dump()
+                )
+            elif isinstance(self._configuration.checkpointer, PostgresCheckpointConfig):
+                self._checkpointer = AsyncPostgresSaver.from_conn_string(
+                    self._configuration.checkpointer.db_url
+                )
+                await self._checkpointer.setup()
+                self._infos["checkpointer"] = (
+                    self._configuration.checkpointer.model_dump()
+                )
             else:
-                raise NotImplementedError("Only SQLite checkpointer is supported.")
+                raise NotImplementedError(
+                    f"Checkpointer type {type(self._configuration.checkpointer)} is not supported."
+                )
 
         if self._configuration.store:
             raise NotImplementedError("Store functionality is not yet implemented.")
@@ -201,14 +282,24 @@ class LanggraphAgent(agent_base.BaseAgent):
         """Loads a StateGraph instance from a specified path."""
         try:
             module_path, graph_variable_name = graph_definition.rsplit(":", 1)
+            if not module_path.endswith(".py"):
+                module_path += ".py"
         except ValueError:
             raise ValueError(
                 "graph_definition must be in the format 'path/to/file.py:variable_name'"
             ) from None
 
+        # Try loading as a file path first
         try:
+            import os
+            print("Current directory: ", os.getcwd()) # TODO remove
             from pathlib import Path
+
             resolved_path = Path(module_path).resolve()
+            # If the file doesn't exist, it might be a python module path
+            if not resolved_path.exists():
+                raise FileNotFoundError
+
             spec = importlib.util.spec_from_file_location(
                 graph_variable_name, str(resolved_path)
             )
@@ -219,17 +310,36 @@ class LanggraphAgent(agent_base.BaseAgent):
             spec.loader.exec_module(module)
 
             graph_builder = getattr(module, graph_variable_name)
-        except (FileNotFoundError, ImportError, AttributeError) as e:
-            raise ValueError(
+            return self._validate_graph_builder(graph_builder, module_path, graph_variable_name)
+
+        except (FileNotFoundError, ImportError):
+            # Fallback: try loading as a python module
+            try:
+                module = importlib.import_module(module_path)
+                graph_builder = getattr(module, graph_variable_name)
+                return self._validate_graph_builder(graph_builder, module_path, graph_variable_name)
+            except ImportError as e:
+                 raise ValueError(
+                    f"Failed to load agent from {graph_definition}. Checked file path and python module: {e}"
+                ) from e
+            except AttributeError as e:
+                raise ValueError(
+                    f"Variable '{graph_variable_name}' not found in module {module_path}: {e}"
+                ) from e
+        except Exception as e:
+             raise ValueError(
                 f"Failed to load agent from {graph_definition}: {e}"
             ) from e
 
-        if not isinstance(graph_builder, StateGraph):
+    def _validate_graph_builder(self, graph_builder: Any, module_path: str, graph_variable_name: str) -> StateGraph:
+         # TODO to remove, dirty fix for template deepagent langgraph
+        if not isinstance(graph_builder, StateGraph) and not isinstance(
+            graph_builder, CompiledStateGraph
+        ):
             raise TypeError(
                 f"The variable '{graph_variable_name}' from {module_path} is not a StateGraph instance."
             )
-
-        return graph_builder
+        return graph_builder # type: ignore[return-value]
 
     async def invoke(self, message: Any) -> Any:
         """Process a single input to chat with the agent.

idun_agent_engine/core/app_factory.py

@@ -8,12 +8,14 @@ setting up routes, dependencies, and lifecycle management behind the scenes.
 from typing import Any
 
 from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 
 from ..server.lifespan import lifespan
 from ..server.routers.agent import agent_router
 from ..server.routers.base import base_router
 from .config_builder import ConfigBuilder
 from .engine_config import EngineConfig
+from .._version import __version__
 
 
 def create_app(
@@ -48,11 +50,19 @@ def create_app(
         lifespan=lifespan,
         title="Idun Agent Engine Server",
         description="A production-ready server for conversational AI agents",
-        version="0.1.0",
+        version=__version__,
         docs_url="/docs",
         redoc_url="/redoc",
     )
 
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
     # Store configuration in app state for lifespan to use
     app.state.engine_config = validated_config