fix: fix type errors and add logs

Author: yaozheng-fang

veadk/__init__.py

@@ -12,17 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .version import VERSION
+from typing import TYPE_CHECKING
+
+from veadk.version import VERSION
+
+if TYPE_CHECKING:
+    from veadk.agent import Agent
+    from veadk.runner import Runner
 
 
 # Lazy loading for `Agent` class
 def __getattr__(name):
     if name == "Agent":
-        from .agent import Agent
+        from veadk.agent import Agent
 
         return Agent
     if name == "Runner":
-        from .runner import Runner
+        from veadk.runner import Runner
 
         return Runner
     raise AttributeError(f"module 'veadk' has no attribute '{name}'")

veadk/agent.py

@@ -17,6 +17,7 @@
 from typing import Optional
 
 from google.adk.agents import LlmAgent, RunConfig
+from google.adk.agents.base_agent import BaseAgent
 from google.adk.agents.llm_agent import ToolUnion
 from google.adk.agents.run_config import StreamingMode
 from google.adk.models.lite_llm import LiteLlm
@@ -39,7 +40,6 @@
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.utils.logger import get_logger
 from veadk.utils.patches import patch_asyncio
-from google.adk.agents.base_agent import BaseAgent
 
 patch_asyncio()
 logger = get_logger(__name__)
@@ -70,9 +70,7 @@ class Agent(LlmAgent):
     model_api_base: str = getenv("MODEL_AGENT_API_BASE", DEFAULT_MODEL_AGENT_API_BASE)
     """The api base of the model for agent running."""
 
-    model_api_key: str = Field(
-        ..., default_factory=lambda: getenv("MODEL_AGENT_API_KEY")
-    )
+    model_api_key: str = Field(default_factory=lambda: getenv("MODEL_AGENT_API_KEY"))
     """The api key of the model for agent running."""
 
     tools: list[ToolUnion] = []
@@ -244,8 +242,13 @@ async def run(
                 user_id=user_id,
                 session_id=session_id,
             )
-            await self.long_term_memory.add_session_to_memory(session)
-            logger.info(f"Add session `{session.id}` to your long-term memory.")
+            if session:
+                await self.long_term_memory.add_session_to_memory(session)
+                logger.info(f"Add session `{session.id}` to your long-term memory.")
+            else:
+                logger.error(
+                    f"Session {session_id} not found in session service, cannot save to long-term memory."
+                )
 
         if collect_runtime_data:
             eval_set_recorder = EvalSetRecorder(session_service, eval_set_id)
@@ -254,6 +257,6 @@ async def run(
 
         if self.tracers:
             for tracer in self.tracers:
-                tracer.dump(user_id, session_id)
+                tracer.dump(user_id=user_id, session_id=session_id)
 
         return final_output

veadk/cli/services/vefaas/vefaas.py

@@ -304,6 +304,7 @@ def find_app_id_by_name(self, name: str):
         for app in apps:
             if app["Name"] == name:
                 return app["Id"]
+        logger.warning(f"Application with name {name} not found.")
         return None
 
     def delete(self, app_id: str):

veadk/cloud/cloud_agent_engine.py

@@ -166,6 +166,10 @@ def remove(self, app_name: str):
             return
         else:
             app_id = self._vefaas_service.find_app_id_by_name(app_name)
+            if not app_id:
+                raise ValueError(
+                    f"Cloud app {app_name} not found, cannot delete it. Please check the app name."
+                )
             self._vefaas_service.delete(app_id)
 
     def update_function_code(

veadk/cloud/cloud_app.py

@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 from typing import Any
 from uuid import uuid4
 
-import json
 import httpx
 from a2a.client import A2ACardResolver, A2AClient
 from a2a.types import AgentCard, Message, MessageSendParams, SendMessageRequest
@@ -82,11 +82,18 @@ def _get_vefaas_endpoint(
         from veadk.cli.services.vefaas.vefaas import VeFaaS
 
         vefaas_client = VeFaaS(access_key=volcengine_ak, secret_key=volcengine_sk)
+
         app = vefaas_client.get_application_details(
             app_id=self.vefaas_application_id,
             app_name=self.vefaas_application_name,
         )
+
+        if not app:
+            raise ValueError(
+                f"VeFaaS CloudAPP with application_id `{self.vefaas_application_id}` or application_name `{self.vefaas_application_name}` not found."
+            )
         cloud_resource = json.loads(app["CloudResource"])
+
         try:
             vefaas_endpoint = cloud_resource["framework"]["url"]["system_url"]
         except Exception as e:
@@ -180,11 +187,19 @@ async def message_send(
                     id=uuid4().hex,
                     params=MessageSendParams(**send_message_payload),
                 )
+
                 res = await a2a_client.send_message(
                     message_send_request,
                     http_kwargs={"timeout": httpx.Timeout(timeout)},
                 )
-                return res.root.result
+
+                logger.debug(
+                    f"Message sent to cloud app {self.vefaas_application_name} with response: {res}"
+                )
+
+                # we ignore type checking here, because the response
+                # from CloudApp will not be `Task` type
+                return res.root.result  # type: ignore
             except Exception as e:
                 # TODO(floritange): show error log on VeFaaS function
                 print(e)
@@ -194,7 +209,7 @@ async def message_send(
 def get_message_id(message: Message):
     """Get the messageId of the a2a message"""
     if getattr(message, "messageId", None):
-        # Compatible with the messageId of the old version
-        return message.messageId
+        # Compatible with the messageId of the old a2a-python version (<0.3.0) in cloud app
+        return message.messageId  # type: ignore
     else:
         return message.message_id

veadk/database/database_adapter.py

@@ -15,8 +15,8 @@
 import re
 import time
 from typing import BinaryIO, TextIO
-from veadk.database.base_database import BaseDatabase
 
+from veadk.database.base_database import BaseDatabase
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -41,7 +41,7 @@ def add(self, data: list[str], index: str):
             )
             raise e
 
-    def query(self, query: str, index: str, top_k: int = 0) -> list[str]:
+    def query(self, query: str, index: str, top_k: int = 0) -> list:
         logger.debug(f"Querying Redis database: index={index} query={query}")
 
         # ignore top_k, as KV search only return one result

veadk/database/kv/redis_database.py

@@ -17,7 +17,7 @@
 from typing import Any
 
 import redis
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, Field
 from typing_extensions import override
 
 from veadk.config import getenv
@@ -30,19 +30,19 @@
 
 class RedisDatabaseConfig(BaseModel):
     host: str = Field(
-        default=getenv("DATABASE_REDIS_HOST"),
+        default_factory=lambda: getenv("DATABASE_REDIS_HOST"),
         description="Redis host",
     )
     port: int = Field(
-        default=getenv("DATABASE_REDIS_PORT"),
+        default_factory=lambda: int(getenv("DATABASE_REDIS_PORT")),
         description="Redis port",
     )
     db: int = Field(
-        default=getenv("DATABASE_REDIS_DB"),
+        default_factory=lambda: int(getenv("DATABASE_REDIS_DB")),
         description="Redis db",
     )
     password: str = Field(
-        default=getenv("DATABASE_REDIS_PASSWORD"),
+        default_factory=lambda: getenv("DATABASE_REDIS_PASSWORD"),
         description="Redis password",
     )
     decode_responses: bool = Field(
@@ -53,7 +53,6 @@ class RedisDatabaseConfig(BaseModel):
 
 class RedisDatabase(BaseModel, BaseDatabase):
     config: RedisDatabaseConfig = Field(default_factory=RedisDatabaseConfig)
-    _client: redis.Redis = PrivateAttr(default=None)
 
     def model_post_init(self, context: Any, /) -> None:
         try:
@@ -64,6 +63,7 @@ def model_post_init(self, context: Any, /) -> None:
                 password=self.config.password,
                 decode_responses=self.config.decode_responses,
             )
+
             self._client.ping()
             logger.info("Connected to Redis successfully.")
         except Exception as e:
@@ -79,10 +79,10 @@ def add(self, key: str, value: str, **kwargs):
             raise e
 
     @override
-    def query(self, key: str, query: str = "", **kwargs) -> list[str]:
+    def query(self, key: str, query: str = "", **kwargs) -> list:
         try:
             result = self._client.lrange(key, 0, -1)
-            return result
+            return result  # type: ignore
         except Exception as e:
             logger.error(f"Failed to search from Redis list key '{key}': {e}")
             raise e
@@ -99,8 +99,11 @@ def delete(self, **kwargs):
 
         try:
             # For simple key deletion
+            # We use sync Redis client to delete the key
+            # so the result will be `int`
             result = self._client.delete(key)
-            if result > 0:
+
+            if result > 0:  # type: ignore
                 logger.info(f"Deleted key `{key}` from Redis.")
             else:
                 logger.info(f"Key `{key}` not found in Redis. Skipping deletion.")

veadk/database/relational/mysql_database.py

@@ -17,7 +17,7 @@
 from typing import Any
 
 import pymysql
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, Field
 from typing_extensions import override
 
 from veadk.config import getenv
@@ -30,32 +30,30 @@
 
 class MysqlDatabaseConfig(BaseModel):
     host: str = Field(
-        default=getenv("DATABASE_MYSQL_HOST"),
+        default_factory=lambda: getenv("DATABASE_MYSQL_HOST"),
         description="Mysql host",
     )
     user: str = Field(
-        default=getenv("DATABASE_MYSQL_USER"),
+        default_factory=lambda: getenv("DATABASE_MYSQL_USER"),
         description="Mysql user",
     )
     password: str = Field(
-        default=getenv("DATABASE_MYSQL_PASSWORD"),
+        default_factory=lambda: getenv("DATABASE_MYSQL_PASSWORD"),
         description="Mysql password",
     )
     database: str = Field(
-        default=getenv("DATABASE_MYSQL_DATABASE"),
+        default_factory=lambda: getenv("DATABASE_MYSQL_DATABASE"),
         description="Mysql database",
     )
     charset: str = Field(
-        default=getenv("DATABASE_MYSQL_CHARSET", "utf8mb4"),
+        default_factory=lambda: getenv("DATABASE_MYSQL_CHARSET", "utf8mb4"),
         description="Mysql charset",
     )
 
 
 class MysqlDatabase(BaseModel, BaseDatabase):
     config: MysqlDatabaseConfig = Field(default_factory=MysqlDatabaseConfig)
 
-    _connection: pymysql.Connection = PrivateAttr(default=None)
-
     def model_post_init(self, context: Any, /) -> None:
         self._connection = pymysql.connect(
             host=self.config.host,
@@ -65,6 +63,9 @@ def model_post_init(self, context: Any, /) -> None:
             charset=self.config.charset,
             cursorclass=pymysql.cursors.DictCursor,
         )
+        self._connection.ping()
+        logger.info("Connected to MySQL successfully.")
+
         self._type = "mysql"
 
     def table_exists(self, table: str) -> bool:
@@ -83,7 +84,7 @@ def add(self, sql: str, params=None, **kwargs):
             self._connection.commit()
 
     @override
-    def query(self, sql: str, params=None, **kwargs) -> list[str]:
+    def query(self, sql: str, params=None, **kwargs) -> tuple[dict[str, Any], ...]:
         with self._connection.cursor() as cursor:
             cursor.execute(sql, params)
             return cursor.fetchall()

veadk/database/vector/opensearch_vector_database.py

@@ -32,23 +32,29 @@
 
 class OpenSearchVectorDatabaseConfig(BaseModel):
     host: str = Field(
-        default=getenv("DATABASE_OPENSEARCH_HOST"),
+        default_factory=lambda: getenv("DATABASE_OPENSEARCH_HOST"),
         description="OpenSearch host",
     )
+
     port: str | int = Field(
-        default=getenv("DATABASE_OPENSEARCH_PORT"),
+        default_factory=lambda: getenv("DATABASE_OPENSEARCH_PORT"),
         description="OpenSearch port",
     )
+
     username: Optional[str] = Field(
-        default=getenv("DATABASE_OPENSEARCH_USERNAME"),
+        default_factory=lambda: getenv("DATABASE_OPENSEARCH_USERNAME"),
         description="OpenSearch username",
     )
+
     password: Optional[str] = Field(
-        default=getenv("DATABASE_OPENSEARCH_PASSWORD"),
+        default_factory=lambda: getenv("DATABASE_OPENSEARCH_PASSWORD"),
         description="OpenSearch password",
     )
+
     secure: bool = Field(default=True, description="Whether enable SSL")
+
     verify_certs: bool = Field(default=False, description="Whether verify SSL certs")
+
     auth_method: Literal["basic", "aws_managed_iam"] = Field(
         default="basic", description="OpenSearch auth method"
     )
@@ -231,15 +237,16 @@ def get_all_docs(self, collection_name: str, size: int = 10000) -> list[dict]:
             for hit in response["hits"]["hits"]
         ]
 
-    def delete_by_query(self, collection_name: str, query: str):
+    def delete_by_query(self, collection_name: str, query: str) -> Any:
         """Delete docs by query in one index of OpenSearch"""
         if not self.collection_exists(collection_name):
             raise ValueError(f"Collection {collection_name} does not exist.")
 
-        query = {"query": {"match": {"page_content": query}}}
+        query_payload = {"query": {"match": {"page_content": query}}}
         response = self._opensearch_client.delete_by_query(
-            index=collection_name, body=query
+            index=collection_name, body=query_payload
         )
+
         self._opensearch_client.indices.refresh(index=collection_name)
         return response
 

veadk/database/viking/viking_database.py

@@ -136,13 +136,17 @@ def _upload_to_tos(
         file_ext = kwargs.get(
             "file_ext", ".pdf"
         )  # when bytes data, file_ext is required
+
         ak = self.config.volcengine_ak
         sk = self.config.volcengine_sk
+
         tos_bucket = self.config.tos.bucket
         tos_endpoint = self.config.tos.endpoint
         tos_region = self.config.tos.region
         tos_key = self.config.tos.base_key
+
         client = tos.TosClientV2(ak, sk, tos_endpoint, tos_region, max_connections=1024)
+
         if isinstance(data, str) and os.path.isfile(data):  # Process file path
             file_ext = os.path.splitext(data)[1]
             new_key = f"{tos_key}/{str(uuid.uuid4())}{file_ext}"