veadk-python 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

veadk/memory/short_term_memory_processor.py

@@ -51,13 +51,16 @@ class ShortTermMemoryProcessor:
         messages = []
         for event in session.events:
             content = event.content
+            if not content or not content.parts:
+                continue
             message = {
                 "role": content.role,
                 "content": content.parts[0].text,
             }
             messages.append(message)
 
-        prompt = render_prompt(messages=str(messages))
+        prompt = render_prompt(messages=messages)
+
         res = completion(
             model=getenv("MODEL_AGENT_PROVIDER") + "/" + getenv("MODEL_AGENT_NAME"),
             base_url=getenv("MODEL_AGENT_API_BASE"),
@@ -69,7 +72,9 @@ class ShortTermMemoryProcessor:
                 }
             ],
         )
-        extracted_messages = json.loads(res.choices[0].message.content)
+        logger.debug(f"Response from memory optimization model: {res}")
+
+        extracted_messages = json.loads(res.choices[0].message.content)  # type: ignore
         logger.debug(f"Abstracted messages: {extracted_messages}")
 
         session.events = []

veadk/prompts/prompt_optimization.py

@@ -112,33 +112,26 @@ def render_prompt_with_jinja2(agent: Agent):
 
     tools = []
     for tool in agent.tools:
+        _tool_type = ""
+        _tools = []
         if isinstance(tool, Callable):
-            type = "function"
-
-            _tool = FunctionTool(tool)
-            tools.append(
-                {
-                    "name": _tool.name,
-                    "description": _tool.description,
-                    "arguments": str(
-                        _tool._get_declaration().model_dump()["parameters"]
-                    ),
-                    "type": type,
-                }
-            )
-        elif isinstance(tool, MCPToolset):
-            type = "tool"
+            _tool_type = "function"
+            _tools = [FunctionTool(tool)]
 
+        elif isinstance(tool, MCPToolset):
+            _tool_type = "tool"
             _tools = asyncio.run(tool.get_tools())
-            for _tool in _tools:
+
+        for _tool in _tools:
+            if _tool and _tool._get_declaration():
                 tools.append(
                     {
                         "name": _tool.name,
                         "description": _tool.description,
                         "arguments": str(
-                            _tool._get_declaration().model_dump()["parameters"]
+                            _tool._get_declaration().model_dump()["parameters"]  # type: ignore
                         ),
-                        "type": type,
+                        "type": _tool_type,
                     }
                 )
 
@@ -148,7 +141,6 @@ def render_prompt_with_jinja2(agent: Agent):
             "name": agent.name,
             "model": agent.model_name,
             "description": agent.description,
-            # "tools": str(agent.tools),
         },
         "tools": tools,
     }

veadk/runner.py

@@ -57,12 +57,14 @@ class Runner:
         self.short_term_memory = short_term_memory
         self.session_service = short_term_memory.session_service
 
+        # prevent VeRemoteAgent has no long-term memory attr
         if isinstance(self.agent, Agent):
             self.long_term_memory = self.agent.long_term_memory
+            for tracer in self.agent.tracers:
+                tracer.set_app_name(self.app_name)
         else:
             self.long_term_memory = None
 
-        # maintain a in-memory runner for fast inference
         self.runner = ADKRunner(
             app_name=self.app_name,
             agent=self.agent,
@@ -70,10 +72,6 @@ class Runner:
             memory_service=self.long_term_memory,
         )
 
-        if getattr(self.agent, "tracers", None):
-            for tracers in self.agent.tracers:
-                tracers.set_app_name(self.app_name)
-
     def _convert_messages(self, messages) -> list:
         if isinstance(messages, str):
             messages = [types.Content(role="user", parts=[types.Part(text=messages)])]
@@ -126,6 +124,7 @@ class Runner:
                         logger.debug(f"Function call: {function_call}")
                 elif (
                     event.content is not None
+                    and event.content.parts
                     and event.content.parts[0].text is not None
                     and len(event.content.parts[0].text.strip()) > 0
                 ):
@@ -146,8 +145,9 @@ class Runner:
         messages: RunnerMessage,
         session_id: str,
         stream: bool = False,
+        save_tracing_data: bool = False,
     ):
-        messages: list = self._convert_messages(messages)
+        converted_messages: list = self._convert_messages(messages)
 
         await self.short_term_memory.create_session(
             app_name=self.app_name, user_id=self.user_id, session_id=session_id
@@ -156,23 +156,32 @@ class Runner:
         logger.info("Begin to process user messages.")
 
         final_output = ""
-        for message in messages:
-            final_output = await self._run(session_id, message, stream)
+        for converted_message in converted_messages:
+            final_output = await self._run(session_id, converted_message, stream)
 
         # try to save tracing file
-        if isinstance(self.agent, Agent):
+        if save_tracing_data:
             self.save_tracing_file(session_id)
 
         return final_output
 
     def save_tracing_file(self, session_id: str) -> str:
+        if not isinstance(self.agent, Agent):
+            logger.warning(
+                (
+                    "The agent is not an instance of VeADK Agent, cannot save tracing file."
+                )
+            )
+            return ""
+
         if not self.agent.tracers:
-            return
+            logger.warning("No tracer is configured in the agent.")
+            return ""
 
         try:
             dump_path = ""
             for tracer in self.agent.tracers:
-                dump_path = tracer.dump(self.user_id, session_id)
+                dump_path = tracer.dump(user_id=self.user_id, session_id=session_id)
 
             return dump_path
         except Exception as e:
@@ -196,61 +205,69 @@ class Runner:
             user_id=self.user_id,
             session_id=session_id,
         )
-        await self.long_term_memory.add_session_to_memory(session)
-        logger.info(f"Add session `{session.id}` to long-term memory.")
-
-    async def run_with_final_event(
-        self,
-        messages: RunnerMessage,
-        session_id: str,
-    ):
-        """non-streaming run with final event"""
-        messages: list = self._convert_messages(messages)
-
-        await self.short_term_memory.create_session(
-            app_name=self.app_name, user_id=self.user_id, session_id=session_id
-        )
-
-        logger.info("Begin to process user messages.")
-
-        final_event = ""
-        async for event in self.runner.run_async(
-            user_id=self.user_id, session_id=session_id, new_message=messages[0]
-        ):
-            if event.get_function_calls():
-                for function_call in event.get_function_calls():
-                    logger.debug(f"Function call: {function_call}")
-            elif (
-                not event.partial
-                and event.content.parts[0].text is not None
-                and len(event.content.parts[0].text.strip()) > 0
-            ):
-                final_event = event.model_dump_json(exclude_none=True, by_alias=True)
-
-        return final_event
-
-    async def run_sse(
-        self,
-        session_id: str,
-        prompt: str,
-    ):
-        message = types.Content(role="user", parts=[types.Part(text=prompt)])
-
-        await self.short_term_memory.create_session(
-            app_name=self.app_name, user_id=self.user_id, session_id=session_id
-        )
-
-        logger.info("Begin to process user messages under SSE method.")
+        if not session:
+            logger.error(
+                f"Session {session_id} not found in session service, cannot save to long-term memory."
+            )
+            return
 
-        async for event in self.runner.run_async(
-            user_id=self.user_id,
-            session_id=session_id,
-            new_message=message,
-            run_config=RunConfig(streaming_mode=StreamingMode.SSE),
-        ):
-            # Format as SSE data
-            sse_event = event.model_dump_json(exclude_none=True, by_alias=True)
-            if event.get_function_calls():
-                for function_call in event.get_function_calls():
-                    logger.debug(f"SSE function call event: {sse_event}")
-            yield f"data: {sse_event}\n\n"
+        await self.long_term_memory.add_session_to_memory(session)
+        logger.info(f"Add session `{session.id}` to long term memory.")
+
+    # [deprecated] we will not host a chat-service in VeADK, so the following two methods are deprecated
+
+    # async def run_with_final_event(
+    #     self,
+    #     messages: RunnerMessage,
+    #     session_id: str,
+    # ):
+    #     """non-streaming run with final event"""
+    #     messages: list = self._convert_messages(messages)
+
+    #     await self.short_term_memory.create_session(
+    #         app_name=self.app_name, user_id=self.user_id, session_id=session_id
+    #     )
+
+    #     logger.info("Begin to process user messages.")
+
+    #     final_event = ""
+    #     async for event in self.runner.run_async(
+    #         user_id=self.user_id, session_id=session_id, new_message=messages[0]
+    #     ):
+    #         if event.get_function_calls():
+    #             for function_call in event.get_function_calls():
+    #                 logger.debug(f"Function call: {function_call}")
+    #         elif (
+    #             not event.partial
+    #             and event.content.parts[0].text is not None
+    #             and len(event.content.parts[0].text.strip()) > 0
+    #         ):
+    #             final_event = event.model_dump_json(exclude_none=True, by_alias=True)
+
+    #     return final_event
+
+    # async def run_sse(
+    #     self,
+    #     session_id: str,
+    #     prompt: str,
+    # ):
+    #     message = types.Content(role="user", parts=[types.Part(text=prompt)])
+
+    #     await self.short_term_memory.create_session(
+    #         app_name=self.app_name, user_id=self.user_id, session_id=session_id
+    #     )
+
+    #     logger.info("Begin to process user messages under SSE method.")
+
+    #     async for event in self.runner.run_async(
+    #         user_id=self.user_id,
+    #         session_id=session_id,
+    #         new_message=message,
+    #         run_config=RunConfig(streaming_mode=StreamingMode.SSE),
+    #     ):
+    #         # Format as SSE data
+    #         sse_event = event.model_dump_json(exclude_none=True, by_alias=True)
+    #         if event.get_function_calls():
+    #             for function_call in event.get_function_calls():
+    #                 logger.debug(f"SSE function call event: {sse_event}")
+    #         yield f"data: {sse_event}\n\n"

veadk/tools/load_knowledgebase_tool.py

@@ -26,7 +26,7 @@ from typing_extensions import override
 from veadk.knowledgebase import KnowledgeBase
 
 if TYPE_CHECKING:
-    from google.adk.models import LlmRequest
+    from google.adk.models.llm_request import LlmRequest
 
 
 knowledgebase: KnowledgeBase | None = None
@@ -85,7 +85,8 @@ async def load_knowledgebase(
     Returns:
       A list of knowledgebase results.
     """
-    search_knowledgebase_response = await tool_context.search_knowledgebase(
+
+    search_knowledgebase_response = await tool_context.search_knowledgebase(  # type: ignore[attr-defined]
         query, tool_context._invocation_context.app_name
     )
     return LoadKnowledgebaseResponse(

veadk/tracing/base_tracer.py

@@ -14,10 +14,11 @@
 
 import json
 from abc import ABC, abstractmethod
-from typing import Optional, Any
+from typing import Any, Optional
 
 from google.adk.agents.callback_context import CallbackContext
-from google.adk.models import LlmRequest, LlmResponse
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
 from google.adk.tools import BaseTool, ToolContext
 from opentelemetry import trace
 
@@ -32,7 +33,7 @@ class BaseTracer(ABC):
         pass
 
     @abstractmethod
-    def dump(self) -> str: ...
+    def dump(self, user_id: str, session_id: str, path: str = "/tmp") -> str: ...
 
     def tracer_hook_before_model(
         self, callback_context: CallbackContext, llm_request: LlmRequest
@@ -107,22 +108,18 @@ class BaseTracer(ABC):
         agent_name = callback_context.agent_name
         attributes["agent.name"] = agent_name
         attributes["app.name"] = app_name
+
         # prompt
         user_content = callback_context.user_content
+        role = None
+        content = None
         if getattr(user_content, "role", None):
             role = getattr(user_content, "role", None)
-        else:
-            role = None
-        if getattr(user_content, "parts", None):
-            content = callback_context.user_content.model_dump(exclude_none=True).get(
-                "parts", None
-            )
-            if content:
-                content = json.dumps(content)
-            else:
-                content = None
-        else:
-            content = None
+
+        if user_content and getattr(user_content, "parts", None):
+            content = user_content.model_dump(exclude_none=True).get("parts", None)
+            content = json.dumps(content) if content else None
+
         if role and content:
             attributes["gen_ai.prompt.0.role"] = role
             attributes["gen_ai.prompt.0.content"] = content
@@ -211,3 +208,18 @@ class BaseTracer(ABC):
 
     def set_app_name(self, app_name):
         self.app_name = app_name
+
+    def do_hooks(self, agent) -> None:
+        if not getattr(agent, "before_model_callback", None):
+            agent.before_model_callback = []
+        if not getattr(agent, "after_model_callback", None):
+            agent.after_model_callback = []
+        if not getattr(agent, "after_tool_callback", None):
+            agent.after_tool_callback = []
+
+        if self.tracer_hook_before_model not in agent.before_model_callback:
+            agent.before_model_callback.append(self.tracer_hook_before_model)
+        if self.tracer_hook_after_model not in agent.after_model_callback:
+            agent.after_model_callback.append(self.tracer_hook_after_model)
+        if self.tracer_hook_after_tool not in agent.after_tool_callback:
+            agent.after_tool_callback.append(self.tracer_hook_after_tool)

veadk/tracing/telemetry/exporters/apmplus_exporter.py

@@ -32,18 +32,15 @@ logger = get_logger(__name__)
 
 class APMPlusExporterConfig(BaseModel):
     endpoint: str = Field(
-        ...,
         default_factory=lambda: getenv(
             "OBSERVABILITY_OPENTELEMETRY_APMPLUS_ENDPOINT",
             "http://apmplus-cn-beijing.volces.com:4317",
         ),
     )
     app_key: str = Field(
-        ...,
         default_factory=lambda: getenv("OBSERVABILITY_OPENTELEMETRY_APMPLUS_API_KEY"),
     )
     service_name: str = Field(
-        ...,
         default_factory=lambda: getenv(
             "OBSERVABILITY_OPENTELEMETRY_APMPLUS_SERVICE_NAME",
             "veadk_tracing_service",

veadk/tracing/telemetry/exporters/base_exporter.py

@@ -26,3 +26,6 @@ class BaseExporter(ABC):
 
     def get_meter_context(self) -> Any:
         pass
+
+    def export(self) -> None:
+        pass

veadk/tracing/telemetry/exporters/cozeloop_exporter.py

@@ -26,20 +26,17 @@ logger = get_logger(__name__)
 
 class CozeloopExporterConfig(BaseModel):
     endpoint: str = Field(
-        ...,
         default_factory=lambda: getenv(
             "OBSERVABILITY_OPENTELEMETRY_COZELOOP_ENDPOINT",
             "https://api.coze.cn/v1/loop/opentelemetry/v1/traces",
         ),
     )
     space_id: str = Field(
-        ...,
         default_factory=lambda: getenv(
             "OBSERVABILITY_OPENTELEMETRY_COZELOOP_SERVICE_NAME"
         ),
     )
     token: str = Field(
-        ...,
         default_factory=lambda: getenv("OBSERVABILITY_OPENTELEMETRY_COZELOOP_API_KEY"),
     )
 

veadk/tracing/telemetry/exporters/inmemory_exporter.py

@@ -20,6 +20,9 @@ from pydantic import BaseModel
 from typing_extensions import override
 
 from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 # ======== Adapted from Google ADK ========
@@ -35,10 +38,16 @@ class _InMemoryExporter(export.SpanExporter):
     @override
     def export(self, spans: typing.Sequence[ReadableSpan]) -> export.SpanExportResult:
         for span in spans:
-            trace_id = span.context.trace_id
-            self.trace_id = trace_id
+            if span.context:
+                trace_id = span.context.trace_id
+                self.trace_id = trace_id
+            else:
+                logger.warning(
+                    f"Span context is missing, failed to get `trace_id`. span: {span}"
+                )
+
             if span.name == "call_llm":
-                attributes = dict(span.attributes)
+                attributes = dict(span.attributes or {})
                 prompt_token = attributes.get("gen_ai.usage.prompt_tokens", None)
                 completion_token = attributes.get(
                     "gen_ai.usage.completion_tokens", None
@@ -48,7 +57,7 @@ class _InMemoryExporter(export.SpanExporter):
                 if completion_token:
                     self.completion_tokens.append(completion_token)
             if span.name == "call_llm":
-                attributes = dict(span.attributes)
+                attributes = dict(span.attributes or {})
                 session_id = attributes.get("gcp.vertex.agent.session_id", None)
                 if session_id:
                     if session_id not in self.session_trace_dict:

veadk/tracing/telemetry/exporters/tls_exporter.py

@@ -26,29 +26,22 @@ logger = get_logger(__name__)
 
 class TLSExporterConfig(BaseModel):
     endpoint: str = Field(
-        ...,
         default_factory=lambda: getenv(
             "OBSERVABILITY_OPENTELEMETRY_TLS_ENDPOINT",
             "https://tls-cn-beijing.volces.com:4318/v1/traces",
         ),
     )
     region: str = Field(
-        ...,
         default_factory=lambda: getenv(
             "OBSERVABILITY_OPENTELEMETRY_TLS_REGION",
             "cn-beijing",
         ),
     )
     topic_id: str = Field(
-        ...,
         default_factory=lambda: getenv("OBSERVABILITY_OPENTELEMETRY_TLS_SERVICE_NAME"),
     )
-    access_key: str = Field(
-        ..., default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY")
-    )
-    secret_key: str = Field(
-        ..., default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY")
-    )
+    access_key: str = Field(default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY"))
+    secret_key: str = Field(default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY"))
 
 
 class TLSExporter(BaseModel, BaseExporter):

veadk/tracing/telemetry/opentelemetry_tracer.py

@@ -20,6 +20,7 @@ from typing import Any
 
 from openinference.instrumentation.google_adk import GoogleADKInstrumentor
 from opentelemetry import trace as trace_api
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk import trace as trace_sdk
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
@@ -28,7 +29,6 @@ from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import override
 
 from veadk.tracing.base_tracer import BaseTracer
-from veadk.tracing.telemetry.exporters.apiserver_exporter import ApiServerExporter
 from veadk.tracing.telemetry.exporters.apmplus_exporter import APMPlusExporter
 from veadk.tracing.telemetry.exporters.base_exporter import BaseExporter
 from veadk.tracing.telemetry.exporters.inmemory_exporter import InMemoryExporter
@@ -96,48 +96,34 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
         GoogleADKInstrumentor().instrument()
 
     def _init_tracer_provider(self) -> None:
-        # 1. get global trace provider
-        global_tracer_provider = trace_api.get_tracer_provider()
-
-        if not isinstance(global_tracer_provider, TracerProvider):
-            logger.info(
-                "Global tracer provider has not been set. Create tracer provider and set it now."
-            )
-            # 1.1 init tracer provider
-            tracer_provider = trace_sdk.TracerProvider()
-            trace_api.set_tracer_provider(tracer_provider)
-            global_tracer_provider = trace_api.get_tracer_provider()
+        # set provider anyway
+        # finally, get global provider
+        tracer_provider = trace_sdk.TracerProvider()
+        trace_api.set_tracer_provider(tracer_provider)
+        global_tracer_provider: TracerProvider = trace_api.get_tracer_provider()  # type: ignore
 
         have_apmplus_exporter = False
-        # 2. check if apmplus exporter is already exist
         for processor in global_tracer_provider._active_span_processor._span_processors:
             if isinstance(processor, (BatchSpanProcessor, SimpleSpanProcessor)):
-                # try to get endpoint, in case of exporter has no _endpoint attribute
-                try:
-                    exporter_endpoint = processor.span_exporter._endpoint
-                    if "apmplus" in exporter_endpoint:
+                if isinstance(processor.span_exporter, OTLPSpanExporter):
+                    if "apmplus" in processor.span_exporter._endpoint:
                         have_apmplus_exporter = True
-                except AttributeError:
-                    # log a warning and pass this exporter
-                    logger.warning(
-                        f"Exporter {processor.span_exporter} has no endpoint."
-                    )
-
-        # 3. add exporters to global tracer_provider
-        # range over a copy of exporters to avoid index issues
+
         if have_apmplus_exporter:
             self.exporters = [
                 e for e in self.exporters if not isinstance(e, APMPlusExporter)
             ]
 
-        for exporter in self.exporters[:]:
+        for exporter in self.exporters:
             processor, resource_attributes = exporter.get_processor()
             if resource_attributes is not None:
                 update_resource_attributions(
                     global_tracer_provider, resource_attributes
                 )
             global_tracer_provider.add_span_processor(processor)
-            logger.debug(f"Add exporter `{exporter.__class__.__name__}` to tracing.")
+            logger.debug(
+                f"Add exporter `{exporter.__class__.__name__}` to OpentelemetryTracer."
+            )
             self._processors.append(processor)
         logger.debug(f"Init OpentelemetryTracer with {len(self.exporters)} exporters.")
 
@@ -198,9 +184,7 @@ class OpentelemetryTracer(BaseModel, BaseTracer):
         self._trace_file_path = file_path
 
         for exporter in self.exporters:
-            if not isinstance(exporter, InMemoryExporter) and not isinstance(
-                exporter, ApiServerExporter
-            ):
+            if not isinstance(exporter, InMemoryExporter):
                 exporter.export()
         logger.info(
             f"OpenTelemetryTracer tracing done, trace id: {self._trace_id} (hex)"

veadk/types.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+
+from veadk.agent import Agent
+from veadk.memory.short_term_memory import ShortTermMemory
 
 
 class MediaMessage(BaseModel):
@@ -21,3 +24,19 @@ class MediaMessage(BaseModel):
 
     media: str
     """Media file (e.g., `.pdf`, `.docx`, `.png`, `.jpg`, `.jpeg`, `.mp4`, `.mp3`, `.wav`, `.txt`) path"""
+
+
+class AgentRunConfig(BaseModel):
+    """Configuration for running an agent on VeFaaS platform."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    app_name: str = Field(
+        default="veadk_vefaas_app", description="The name of the application"
+    )
+
+    agent: Agent = Field(..., description="The root agent instance")
+
+    short_term_memory: ShortTermMemory = Field(
+        default_factory=ShortTermMemory, description="The short-term memory instance"
+    )

veadk/utils/mcp_utils.py

@@ -14,7 +14,7 @@
 
 from typing import Any
 
-from google.adk.tools.mcp_tool.mcp_toolset import (
+from google.adk.tools.mcp_tool.mcp_session_manager import (
     SseConnectionParams,
     StreamableHTTPConnectionParams,
 )