veadk-python 0.2.6__tar.gz → 0.2.8__tar.gz

{veadk_python-0.2.6 → veadk_python-0.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: veadk-python
-Version: 0.2.6
+Version: 0.2.8
 Summary: Volcengine agent development kit, integrations with Volcengine cloud services.
 Author-email: Yaozheng Fang <fangyozheng@gmail.com>, Guodong Li <cu.eric.lee@gmail.com>, Zhi Han <sliverydayday@gmail.com>, Meng Wang <mengwangwm@gmail.com>
 License:                                  Apache License
@@ -218,13 +218,26 @@ Requires-Dist: opentelemetry-exporter-otlp>=1.35.0
 Requires-Dist: opentelemetry-instrumentation-logging>=0.56b0
 Requires-Dist: wrapt>=1.17.2
 Requires-Dist: openai<1.100
-Requires-Dist: volcengine-python-sdk==4.0.3
+Requires-Dist: volcengine-python-sdk>=4.0.3
+Requires-Dist: volcengine>=1.0.193
 Requires-Dist: agent-pilot-sdk>=0.0.9
 Requires-Dist: fastmcp>=2.11.3
 Requires-Dist: cookiecutter>=2.6.0
+Requires-Dist: omegaconf>=2.3.0
+Requires-Dist: llama-index>=0.14.0
+Requires-Dist: llama-index-embeddings-openai-like>=0.2.2
+Requires-Dist: llama-index-llms-openai-like>=0.5.1
+Requires-Dist: llama-index-vector-stores-opensearch>=0.6.1
+Requires-Dist: psycopg2-binary>=2.9.10
+Requires-Dist: pymysql>=1.1.1
 Requires-Dist: opensearch-py==2.8.0
+Provides-Extra: extensions
+Requires-Dist: redis>=5.0; extra == "extensions"
+Requires-Dist: tos>=2.8.4; extra == "extensions"
+Requires-Dist: llama-index-vector-stores-redis>=0.6.1; extra == "extensions"
+Requires-Dist: mcp-server-vikingdb-memory; extra == "extensions"
 Provides-Extra: database
-Requires-Dist: redis>=6.2.0; extra == "database"
+Requires-Dist: redis>=5.0; extra == "database"
 Requires-Dist: pymysql>=1.1.1; extra == "database"
 Requires-Dist: volcengine>=1.0.193; extra == "database"
 Requires-Dist: tos>=2.8.4; extra == "database"

{veadk_python-0.2.6 → veadk_python-0.2.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "veadk-python"
-version = "0.2.6"
+version = "0.2.8"
 description = "Volcengine agent development kit, integrations with Volcengine cloud services."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -22,19 +22,33 @@ dependencies = [
     "opentelemetry-instrumentation-logging>=0.56b0",
     "wrapt>=1.17.2", # For patching built-in functions
     "openai<1.100", # For fix https://github.com/BerriAI/litellm/issues/13710
-    "volcengine-python-sdk==4.0.3", # For Volcengine API
-    "agent-pilot-sdk>=0.0.9",    # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
-    "fastmcp>=2.11.3",    # For running MCP
+    "volcengine-python-sdk>=4.0.3", # For Volcengine API
+    "volcengine>=1.0.193",          # For Volcengine sign
+    "agent-pilot-sdk>=0.0.9", # Prompt optimization by Volcengine AgentPilot/PromptPilot toolkits
+    "fastmcp>=2.11.3", # For running MCP
     "cookiecutter>=2.6.0", # For cloud deploy
-    "opensearch-py==2.8.0"    # For OpenSearch database
+    "omegaconf>=2.3.0", # For agent builder
+    "llama-index>=0.14.0",
+    "llama-index-embeddings-openai-like>=0.2.2",
+    "llama-index-llms-openai-like>=0.5.1",
+    "llama-index-vector-stores-opensearch>=0.6.1",
+    "psycopg2-binary>=2.9.10", # For PostgreSQL database (short term memory)
+    "pymysql>=1.1.1", # For MySQL database (short term memory)
+    "opensearch-py==2.8.0",
 ]
 
 [project.scripts]
 veadk = "veadk.cli.cli:veadk"
 
 [project.optional-dependencies]
+extensions = [
+    "redis>=5.0",                 # For Redis database
+    "tos>=2.8.4",                   # For TOS storage and Viking DB
+    "llama-index-vector-stores-redis>=0.6.1",
+    "mcp-server-vikingdb-memory",
+]
 database = [
-    "redis>=6.2.0",                 # For Redis database
+    "redis>=5.0",                 # For Redis database
     "pymysql>=1.1.1",               # For MySQL database
     "volcengine>=1.0.193",          # For Viking DB
     "tos>=2.8.4",                   # For TOS storage and Viking DB
@@ -77,3 +91,6 @@ exclude = [
     "veadk/integrations/ve_faas/template/*",
     "veadk/integrations/ve_faas/web_template/*"
 ]
+
+[tool.uv.sources]
+mcp-server-vikingdb-memory = { git = "https://github.com/volcengine/mcp-server", subdirectory = "server/mcp_server_vikingdb_memory" }

veadk_python-0.2.8/tests/test_agent.py

@@ -0,0 +1,211 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import Mock, patch
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.models.lite_llm import LiteLlm
+from google.adk.tools import load_memory
+
+from veadk import Agent
+from veadk.consts import (
+    DEFAULT_AGENT_NAME,
+    DEFAULT_MODEL_AGENT_API_BASE,
+    DEFAULT_MODEL_AGENT_NAME,
+    DEFAULT_MODEL_AGENT_PROVIDER,
+    DEFAULT_MODEL_EXTRA_CONFIG,
+)
+from veadk.knowledgebase import KnowledgeBase
+from veadk.memory.long_term_memory import LongTermMemory
+from veadk.tools import load_knowledgebase_tool
+from veadk.tracing.telemetry.opentelemetry_tracer import OpentelemetryTracer
+
+
+def test_agent():
+    knowledgebase = KnowledgeBase(
+        index="test_index",
+        backend="local",
+        backend_config={"embedding_config": {"api_key": "test"}},
+    )
+    long_term_memory = LongTermMemory(backend="local")
+    tracer = OpentelemetryTracer()
+
+    extra_config = {
+        "extra_headers": {"thinking": "test"},
+        "extra_body": {"content": "test"},
+    }
+
+    agent = Agent(
+        model_name="test_model_name",
+        model_provider="test_model_provider",
+        model_api_key="test_model_api_key",
+        model_api_base="test_model_api_base",
+        model_extra_config=extra_config,
+        tools=[],
+        sub_agents=[],
+        knowledgebase=knowledgebase,
+        long_term_memory=long_term_memory,
+        tracers=[tracer],
+    )
+
+    assert agent.model.model == f"{agent.model_provider}/{agent.model_name}"  # type: ignore
+
+    expected_config = DEFAULT_MODEL_EXTRA_CONFIG.copy()
+    expected_config["extra_headers"] |= extra_config["extra_headers"]
+    expected_config["extra_body"] |= extra_config["extra_body"]
+
+    assert agent.model_extra_config == expected_config
+
+    assert agent.knowledgebase == knowledgebase
+    assert agent.knowledgebase.backend == "local"  # type: ignore
+    assert load_knowledgebase_tool.knowledgebase == agent.knowledgebase
+    assert load_knowledgebase_tool.load_knowledgebase_tool in agent.tools
+
+    assert agent.long_term_memory.backend == "local"  # type: ignore
+    assert load_memory in agent.tools
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_default_values():
+    agent = Agent()
+
+    assert agent.name == DEFAULT_AGENT_NAME
+
+    assert agent.model_name == DEFAULT_MODEL_AGENT_NAME
+    assert agent.model_provider == DEFAULT_MODEL_AGENT_PROVIDER
+    assert agent.model_api_base == DEFAULT_MODEL_AGENT_API_BASE
+
+    assert agent.tools == []
+    assert agent.sub_agents == []
+    assert agent.knowledgebase is None
+    assert agent.long_term_memory is None
+    assert agent.tracers == []
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_without_knowledgebase():
+    agent = Agent()
+
+    assert agent.knowledgebase is None
+    assert load_knowledgebase_tool.load_knowledgebase_tool not in agent.tools
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_without_long_term_memory():
+    agent = Agent()
+
+    assert agent.long_term_memory is None
+    assert load_memory not in agent.tools
+
+
+@patch("veadk.agent.LiteLlm")
+def test_agent_model_creation(mock_lite_llm):
+    mock_model = Mock()
+    mock_lite_llm.return_value = mock_model
+
+    agent = Agent(
+        model_name="test_model",
+        model_provider="test_provider",
+        model_api_key="test_key",
+        model_api_base="test_base",
+    )
+
+    mock_lite_llm.assert_called_once()
+    assert agent.model == mock_model
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_with_existing_model():
+    existing_model = LiteLlm(model="test_model")
+    agent = Agent(model=existing_model)
+
+    assert agent.model == existing_model
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_model_extra_config_merge():
+    user_config = {
+        "extra_headers": {"custom": "header"},
+        "extra_body": {"custom": "body"},
+        "other_param": "value",
+    }
+
+    agent = Agent(model_extra_config=user_config)
+
+    expected_headers = DEFAULT_MODEL_EXTRA_CONFIG["extra_headers"].copy()
+    expected_headers["custom"] = "header"
+
+    expected_body = DEFAULT_MODEL_EXTRA_CONFIG["extra_body"].copy()
+    expected_body["custom"] = "body"
+
+    assert agent.model_extra_config["extra_headers"] == expected_headers
+    assert agent.model_extra_config["extra_body"] == expected_body
+    assert agent.model_extra_config["other_param"] == "value"
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_empty_model_extra_config():
+    agent = Agent(model_extra_config={})
+
+    assert (
+        agent.model_extra_config["extra_headers"]
+        == DEFAULT_MODEL_EXTRA_CONFIG["extra_headers"]
+    )
+    assert (
+        agent.model_extra_config["extra_body"]
+        == DEFAULT_MODEL_EXTRA_CONFIG["extra_body"]
+    )
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_with_tools():
+    mock_tool = Mock()
+    agent = Agent(tools=[mock_tool])
+
+    assert mock_tool in agent.tools
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_with_sub_agents():
+    adk_agent = LlmAgent(name="agent")
+    veadk_agent = Agent(name="agent")
+    agent = Agent(sub_agents=[adk_agent, veadk_agent])
+
+    assert adk_agent in agent.sub_agents
+    assert veadk_agent in agent.sub_agents
+    assert adk_agent.parent_agent == agent
+    assert veadk_agent.parent_agent == agent
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_with_tracers():
+    tracer1 = OpentelemetryTracer()
+    tracer2 = OpentelemetryTracer()
+
+    agent = Agent(tracers=[tracer1, tracer2])
+
+    assert len(agent.tracers) == 2
+    assert tracer1 in agent.tracers
+    assert tracer2 in agent.tracers
+
+
+@patch.dict("os.environ", {"MODEL_AGENT_API_KEY": "mock_api_key"})
+def test_agent_custom_name_and_description():
+    custom_name = "CustomAgent"
+    custom_description = "A custom agent for testing"
+
+    agent = Agent(name=custom_name, description=custom_description)
+
+    assert agent.name == custom_name
+    assert agent.description == custom_description

{veadk_python-0.2.6 → veadk_python-0.2.8}/tests/test_evaluator.py

@@ -62,60 +62,61 @@ EVAL_SET_DATA = {
 TRACE_SET_DATA = [
     {
         "name": "execute_tool get_city_weather",
-        "span_id": 5421848634094108689,
-        "trace_id": 115143748123782151752771111946932434777,
-        "start_time": 1754884672226444000,
-        "end_time": 1754884672226993000,
+        "span_id": 4497348974122733469,
+        "trace_id": 142655176138954930885272077198014871976,
+        "start_time": 1758158957162250000,
+        "end_time": 1758158957162426000,
         "attributes": {
             "gen_ai.tool.name": "get_city_weather",
-            "gen_ai.tool.description": "Retrieves the weather information of a given city. the args must in English",
-            "gen_ai.tool.call.id": "call_6ow5622pvcouw3tpvr7rfqtl",
-            "gcp.vertex.agent.tool_call_args": '{"city": "Xi\'an"}',
+            "gen_ai.tool.input": '{"name": "get_city_weather", "description": "Retrieves the weather information of a given city. the args must in English", "parameters": {"city": "Beijing"}}',
+            "gen_ai.tool.output": '{"id": "call_w4bj25flpvs74zgyyiquqh5s", "name": "get_city_weather", "response": {"result": "Sunny, 25°C"}}',
         },
-        "parent_span_id": 7997784243558253239,
+        "parent_span_id": 574819447039686650,
     },
     {
         "name": "call_llm",
-        "span_id": 7997784243558253239,
-        "trace_id": 115143748123782151752771111946932434777,
+        "span_id": 574819447039686650,
+        "trace_id": 142655176138954930885272077198014871976,
+        "start_time": 1758158945807630000,
+        "end_time": 1758158957171304000,
         "attributes": {
-            "session.id": "veadk_example_session",
-            "user.id": "veadk_default_user",
+            "gen_ai.app.name": "veadk_default_app",
+            "gen_ai.user.id": "veadk_default_user",
+            "gen_ai.prompt.0.role": "user",
+            "gen_ai.prompt.0.content": "How is the weather like in BeiJing?",
         },
-        "parent_span_id": 14844888006539887900,
+        "parent_span_id": 13789664766018020416,
     },
     {
         "name": "call_llm",
-        "span_id": 7789424022423491416,
-        "trace_id": 115143748123782151752771111946932434777,
+        "span_id": 9007934154052797946,
+        "trace_id": 142655176138954930885272077198014871976,
+        "start_time": 1758158957171713000,
+        "end_time": 1758158964035230000,
         "attributes": {
-            "session.id": "veadk_example_session",
-            "user.id": "veadk_default_user",
+            "gen_ai.app.name": "veadk_default_app",
+            "gen_ai.user.id": "veadk_default_user",
+            "gen_ai.prompt.0.content": "How is the weather like in BeiJing?",
+            "gen_ai.completion.0.content": "The weather in Beijing is sunny with a temperature of 25°C.",
         },
-        "parent_span_id": 14844888006539887900,
+        "parent_span_id": 13789664766018020416,
     },
     {
         "name": "agent_run [chat_robot]",
-        "span_id": 14844888006539887900,
-        "trace_id": 115143748123782151752771111946932434777,
-        "attributes": {
-            "session.id": "veadk_example_session",
-            "user.id": "veadk_default_user",
-        },
-        "parent_span_id": 2943363177785645047,
+        "span_id": 13789664766018020416,
+        "trace_id": 142655176138954930885272077198014871976,
+        "start_time": 1758158945807350000,
+        "end_time": 1758158964035291000,
+        "attributes": {},
+        "parent_span_id": 5589459087402275636,
     },
     {
-        "name": "invocation [veadk_default_app]",
-        "span_id": 2943363177785645047,
-        "trace_id": 115143748123782151752771111946932434777,
-        "start_time": 1754884660687962000,
-        "end_time": 1754884676664833000,
-        "attributes": {
-            "input.value": '{"user_id": "veadk_default_user", "session_id": "veadk_example_session", "new_message": "{\\"parts\\": [{\\"video_metadata\\": null, \\"thought\\": null, \\"inline_data\\": null, \\"file_data\\": null, \\"thought_signature\\": null, \\"code_execution_result\\": null, \\"executable_code\\": null, \\"function_call\\": null, \\"function_response\\": null, \\"text\\": \\"How is the weather like in Xi\'an?\\"}], \\"role\\": \\"user\\"}", "run_config": "{\\"speech_config\\": null, \\"response_modalities\\": null, \\"save_input_blobs_as_artifacts\\": false, \\"support_cfc\\": false, \\"streaming_mode\\": \\"StreamingMode.NONE\\", \\"output_audio_transcription\\": null, \\"input_audio_transcription\\": null, \\"realtime_input_config\\": null, \\"enable_affective_dialog\\": null, \\"proactivity\\": null, \\"max_llm_calls\\": 500}"}',
-            "user.id": "veadk_default_user",
-            "session.id": "veadk_example_session",
-            "output.value": '{"content":{"parts":[{"text":"The weather in Xi\'an is cool, with a temperature of 18\u00b0C."}],"role":"model"},"partial":false,"usage_metadata":{"candidates_token_count":132,"prompt_token_count":547,"total_token_count":679},"invocation_id":"e-ea6bb35b-c3f0-4c5c-b127-c71c7d6d6441","author":"chat_robot","actions":{"state_delta":{},"artifact_delta":{},"requested_auth_configs":{}},"id":"c0929124-9be0-4f75-a6ba-f7a531c9ccb6","timestamp":1754884672.227546}',
-        },
+        "name": "invocation",
+        "span_id": 5589459087402275636,
+        "trace_id": 142655176138954930885272077198014871976,
+        "start_time": 1758158945807233000,
+        "end_time": 1758158964035304000,
+        "attributes": {},
         "parent_span_id": None,
     },
 ]
@@ -154,8 +155,8 @@ def test_tracing_file_to_evalset():
     assert len(base_evaluator.invocation_list) == 1
     assert len(base_evaluator.invocation_list[0].invocations) == 1
     assert (
-        base_evaluator.invocation_list[0].invocations[0].invocation_id
-        == "e-ea6bb35b-c3f0-4c5c-b127-c71c7d6d6441"
+        base_evaluator.invocation_list[0].invocations[0].expected_output
+        == "The weather in Beijing is sunny with a temperature of 25°C."
     )
 
     os.remove(tracing_file_path)

{veadk_python-0.2.6 → veadk_python-0.2.8}/tests/test_knowledgebase.py

@@ -12,24 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import pytest
 
 from veadk.knowledgebase import KnowledgeBase
+from veadk.knowledgebase.backends.in_memory_backend import InMemoryKnowledgeBackend
 
 
 @pytest.mark.asyncio
 async def test_knowledgebase():
     app_name = "kb_test_app"
-    key = "Supercalifragilisticexpialidocious"
-    kb = KnowledgeBase(backend="local")
-    # Attempt to delete any existing data for the app_name before adding new data
-    kb.add(
-        data=[f"knowledgebase_id is {key}"],
-        app_name=app_name,
-    )
-    res_list = kb.search(
-        query="knowledgebase_id",
+    kb = KnowledgeBase(
+        backend="local",
         app_name=app_name,
+        backend_config={"embedding_config": {"api_key": "test"}},
     )
-    res = "".join(res_list)
-    assert key in res, f"Test failed for backend local res is {res}"
+
+    assert isinstance(kb._backend, InMemoryKnowledgeBackend)

{veadk_python-0.2.6 → veadk_python-0.2.8}/tests/test_long_term_memory.py

@@ -12,11 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import pytest
-from google.adk.events import Event
-from google.adk.sessions import Session
 from google.adk.tools import load_memory
-from google.genai import types
 
 from veadk.agent import Agent
 from veadk.memory.long_term_memory import LongTermMemory
@@ -27,7 +25,11 @@ user_id = "test_user"
 
 @pytest.mark.asyncio
 async def test_long_term_memory():
-    long_term_memory = LongTermMemory(backend="local")
+    long_term_memory = LongTermMemory(
+        backend="local",
+        # app_name=app_name,
+        # user_id=user_id,
+    )
     agent = Agent(
         name="all_name",
         model_name="test_model_name",
@@ -41,31 +43,8 @@ async def test_long_term_memory():
 
     assert load_memory in agent.tools, "load_memory tool not found in agent tools"
 
-    # mock session
-    session = Session(
-        id="test_session_id",
-        app_name=app_name,
-        user_id=user_id,
-        events=[
-            Event(
-                invocation_id="test_invocation_id",
-                author="user",
-                branch=None,
-                content=types.Content(
-                    parts=[types.Part(text="My name is Alice.")],
-                    role="user",
-                ),
-            )
-        ],
-    )
-
-    await long_term_memory.add_session_to_memory(session)
+    assert not agent.long_term_memory._backend
 
-    memories = await long_term_memory.search_memory(
-        app_name=app_name,
-        user_id=user_id,
-        query="Alice",
-    )
-    assert (
-        "Alice" in memories.model_dump()["memories"][0]["content"]["parts"][0]["text"]
-    )
+    # assert agent.long_term_memory._backend.index == build_long_term_memory_index(
+    #     app_name, user_id
+    # )

{veadk_python-0.2.6 → veadk_python-0.2.8}/tests/test_runner.py

@@ -20,7 +20,13 @@ from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.runner import Runner
 
 
+# Import the standalone function instead of accessing as class method
+from veadk.runner import _convert_messages
+
+
 def _test_convert_messages(runner):
+    """Test message conversion logic using standalone _convert_messages function"""
+    # Test single text message conversion
     message = "test message"
     expected_message = [
         types.Content(
@@ -28,11 +34,16 @@ def _test_convert_messages(runner):
             role="user",
         )
     ]
-    actual_message = runner._convert_messages(
-        message, session_id="test_session_id", upload_inline_data_to_tos=True
+    # Modified: Call _convert_messages directly (not as runner method)
+    actual_message = _convert_messages(
+        message,
+        app_name=runner.app_name,
+        user_id=runner.user_id,
+        session_id="test_session_id",
     )
     assert actual_message == expected_message
 
+    # Test multiple text messages conversion
     message = ["test message 1", "test message 2"]
     expected_message = [
         types.Content(
@@ -44,13 +55,18 @@ def _test_convert_messages(runner):
             role="user",
         ),
     ]
-    actual_message = runner._convert_messages(
-        message, session_id="test_session_id", upload_inline_data_to_tos=True
+    # Modified: Call _convert_messages directly (not as runner method)
+    actual_message = _convert_messages(
+        message,
+        app_name=runner.app_name,
+        user_id=runner.user_id,
+        session_id="test_session_id",
     )
     assert actual_message == expected_message
 
 
 def test_runner():
+    """Test Runner class initialization and core properties"""
     short_term_memory = ShortTermMemory()
     long_term_memory = LongTermMemory(backend="local")
     agent = Agent(
@@ -64,10 +80,9 @@ def test_runner():
     runner = Runner(agent=agent, short_term_memory=short_term_memory)
     assert runner.long_term_memory == agent.long_term_memory
 
-    adk_runner = runner.runner
-    assert adk_runner.memory_service == agent.long_term_memory
-    assert adk_runner.session_service == runner.short_term_memory.session_service
+    # Verify inherited ADKRunner properties
+    assert runner.memory_service == agent.long_term_memory
+    assert runner.session_service == runner.short_term_memory.session_service
 
-    _test_convert_messages(runner)
-    _test_convert_messages(runner)
+    # Run message conversion tests
     _test_convert_messages(runner)

{veadk_python-0.2.6 → veadk_python-0.2.8}/tests/test_short_term_memory.py

@@ -15,7 +15,6 @@
 import asyncio
 import os
 
-import veadk.memory.short_term_memory
 from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.utils.misc import formatted_timestamp
 
@@ -35,11 +34,11 @@ def test_short_term_memory():
     )
     assert session is not None
 
-    # database - local
-    veadk.memory.short_term_memory.DEFAULT_LOCAL_DATABASE_PATH = (
-        f"/tmp/tmp_for_test_{formatted_timestamp()}.db"
+    # sqlite
+    memory = ShortTermMemory(
+        backend="sqlite",
+        local_database_path=f"/tmp/tmp_for_test_{formatted_timestamp()}.db",
     )
-    memory = ShortTermMemory(backend="database")
     asyncio.run(
         memory.session_service.create_session(
             app_name="app", user_id="user", session_id="session"
@@ -51,5 +50,5 @@ def test_short_term_memory():
         )
     )
     assert session is not None
-    assert os.path.exists(veadk.memory.short_term_memory.DEFAULT_LOCAL_DATABASE_PATH)
-    os.remove(veadk.memory.short_term_memory.DEFAULT_LOCAL_DATABASE_PATH)
+    assert os.path.exists(memory.local_database_path)
+    os.remove(memory.local_database_path)