fix type error in some files

Author: yaozheng-fang

veadk/database/viking/viking_database.py

@@ -16,7 +16,7 @@
 import json
 import os
 import uuid
-from typing import Any, BinaryIO, Literal, Optional, TextIO
+from typing import Any, BinaryIO, Literal, TextIO
 
 import requests
 import tos
@@ -44,42 +44,44 @@
 
 
 class VolcengineTOSConfig(BaseModel):
-    endpoint: Optional[str] = Field(
-        default=getenv("DATABASE_TOS_ENDPOINT", "tos-cn-beijing.volces.com"),
+    endpoint: str = Field(
+        default_factory=lambda: getenv(
+            "DATABASE_TOS_ENDPOINT", "tos-cn-beijing.volces.com"
+        ),
         description="VikingDB TOS endpoint",
     )
-    region: Optional[str] = Field(
-        default=getenv("DATABASE_TOS_REGION", "cn-beijing"),
+    region: str = Field(
+        default_factory=lambda: getenv("DATABASE_TOS_REGION", "cn-beijing"),
         description="VikingDB TOS region",
     )
-    bucket: Optional[str] = Field(
-        default=getenv("DATABASE_TOS_BUCKET"),
+    bucket: str = Field(
+        default_factory=lambda: getenv("DATABASE_TOS_BUCKET"),
         description="VikingDB TOS bucket",
     )
-    base_key: Optional[str] = Field(
+    base_key: str = Field(
         default="veadk",
         description="VikingDB TOS base key",
     )
 
 
 class VikingDatabaseConfig(BaseModel):
-    volcengine_ak: Optional[str] = Field(
-        default=getenv("VOLCENGINE_ACCESS_KEY"),
+    volcengine_ak: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY"),
         description="VikingDB access key",
     )
-    volcengine_sk: Optional[str] = Field(
-        default=getenv("VOLCENGINE_SECRET_KEY"),
+    volcengine_sk: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY"),
         description="VikingDB secret key",
     )
-    project: Optional[str] = Field(
-        default=getenv("DATABASE_VIKING_PROJECT"),
+    project: str = Field(
+        default_factory=lambda: getenv("DATABASE_VIKING_PROJECT"),
         description="VikingDB project name",
     )
-    region: Optional[str] = Field(
-        default=getenv("DATABASE_VIKING_REGION"),
+    region: str = Field(
+        default_factory=lambda: getenv("DATABASE_VIKING_REGION"),
         description="VikingDB region",
     )
-    tos: Optional[VolcengineTOSConfig] = Field(
+    tos: VolcengineTOSConfig = Field(
         default_factory=VolcengineTOSConfig,
         description="VikingDB TOS configuration",
     )

veadk/database/viking/viking_memory_db.py

@@ -18,7 +18,7 @@
 import threading
 import time
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 
 from pydantic import BaseModel, Field
 from volcengine.ApiInfo import ApiInfo
@@ -35,20 +35,20 @@
 
 
 class VikingMemConfig(BaseModel):
-    volcengine_ak: Optional[str] = Field(
-        default=getenv("VOLCENGINE_ACCESS_KEY"),
+    volcengine_ak: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY"),
         description="VikingDB access key",
     )
-    volcengine_sk: Optional[str] = Field(
-        default=getenv("VOLCENGINE_SECRET_KEY"),
+    volcengine_sk: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY"),
         description="VikingDB secret key",
     )
-    project: Optional[str] = Field(
-        default=getenv("DATABASE_VIKING_PROJECT"),
+    project: str = Field(
+        default_factory=lambda: getenv("DATABASE_VIKING_PROJECT"),
         description="VikingDB project name",
     )
-    region: Optional[str] = Field(
-        default=getenv("DATABASE_VIKING_REGION"),
+    region: str = Field(
+        default_factory=lambda: getenv("DATABASE_VIKING_REGION"),
         description="VikingDB region",
     )
 

veadk/evaluation/adk_evaluator/adk_evaluator.py

@@ -19,6 +19,7 @@
 from typing import Any, Optional
 
 from google.adk import Runner
+from google.adk.agents.base_agent import BaseAgent
 from google.adk.artifacts import BaseArtifactService, InMemoryArtifactService
 from google.adk.evaluation.agent_evaluator import (
     NUM_RUNS,
@@ -68,7 +69,7 @@ async def ve_generate_responses(  # done
         eval_set: EvalSet,
         agent: Agent,
         repeat_num: int = 3,
-        agent_name: str = None,
+        agent_name: str | None = None,
     ):
         results = []
 
@@ -90,7 +91,7 @@ async def ve_generate_responses(  # done
     @staticmethod
     async def _ve_generate_inferences_from_root_agent(
         invocations: list[Invocation],
-        root_agent: Agent,
+        root_agent: BaseAgent,
         reset_func: Any,
         initial_session: Optional[SessionInput] = None,
         session_id: Optional[str] = None,
@@ -117,21 +118,15 @@ async def _ve_generate_inferences_from_root_agent(
         if not artifact_service:
             artifact_service = InMemoryArtifactService()
 
-        if getattr(root_agent, "long_term_memory", None) is not None:
-            runner = Runner(
-                app_name=app_name,
-                agent=root_agent,
-                artifact_service=artifact_service,
-                session_service=session_service,
-                memory_service=root_agent.long_term_memory,  # add long_term_memory
-            )
-        else:
-            runner = Runner(
-                app_name=app_name,
-                agent=root_agent,
-                artifact_service=artifact_service,
-                session_service=session_service,
-            )
+        runner = Runner(
+            app_name=app_name,
+            agent=root_agent,
+            artifact_service=artifact_service,
+            session_service=session_service,
+            memory_service=root_agent.long_term_memory
+            if isinstance(root_agent, Agent)
+            else None,
+        )
 
         # Reset agent state for each query
         if callable(reset_func):

veadk/tools/load_knowledgebase_tool.py

@@ -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/telemetry/exporters/inmemory_exporter.py

@@ -47,7 +47,7 @@ def export(self, spans: typing.Sequence[ReadableSpan]) -> export.SpanExportResul
                 )
 
             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
@@ -57,7 +57,7 @@ def export(self, spans: typing.Sequence[ReadableSpan]) -> export.SpanExportResul
                 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: