fix bugs

Author: yaozheng-fang

veadk/configs/model_configs.py

@@ -51,7 +51,7 @@ class EmbeddingModelConfig(BaseSettings):
     dim: int = 2560
     """Embedding dim is different from different models."""
 
-    api_base: str = "https://ark.cn-beijing.volces.com/api/v3/embeddings"
+    api_base: str = "https://ark.cn-beijing.volces.com/api/v3/"
     """The api base of the model for embedding."""
 
     @cached_property

veadk/knowledgebase/backends/opensearch_backend.py

@@ -99,7 +99,12 @@ def add_from_text(self, text: str | list[str]) -> bool:
 
     @override
     def search(self, query: str, top_k: int = 5) -> list[str]:
-        retrieved_nodes = self._retriever.retrieve(query, top_k=top_k)
+        _original_top_k = self._retriever.similarity_top_k  # type: ignore
+        self._retriever.similarity_top_k = top_k  # type: ignore
+
+        retrieved_nodes = self._retriever.retrieve(query)
+
+        self._retriever.similarity_top_k = _original_top_k  # type: ignore
         return [node.text for node in retrieved_nodes]
 
     def _split_documents(self, documents: list[Document]) -> list[BaseNode]:

veadk/knowledgebase/backends/redis_backend.py

@@ -39,7 +39,7 @@ class RedisKnowledgeBackend(BaseKnowledgebaseBackend):
     redis_config: RedisConfig = Field(default_factory=RedisConfig)
     """Redis client configs"""
 
-    embedding_config: EmbeddingModelConfig
+    embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
     """Embedding model configs"""
 
     def model_post_init(self, __context: Any) -> None:
@@ -105,7 +105,12 @@ def add_from_text(self, text: str | list[str]) -> bool:
 
     @override
     def search(self, query: str, top_k: int = 5) -> list[str]:
+        _original_top_k = self._retriever.similarity_top_k  # type: ignore
+        self._retriever.similarity_top_k = top_k  # type: ignore
+
         retrieved_nodes = self._retriever.retrieve(query, top_k=top_k)
+
+        self._retriever.similarity_top_k = _original_top_k  # type: ignore
         return [node.text for node in retrieved_nodes]
 
     def _split_documents(self, documents: list[Document]) -> list[BaseNode]:

veadk/memory/__init__.py

@@ -11,3 +11,25 @@
 # 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 typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from veadk.memory.long_term_memory import LongTermMemory
+    from veadk.memory.short_term_memory import ShortTermMemory
+
+
+# Lazy loading for classes
+def __getattr__(name):
+    if name == "ShortTermMemory":
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        return ShortTermMemory
+    if name == "LongTeremMemory":
+        from veadk.memory.long_term_memory import LongTermMemory
+
+        return LongTermMemory
+    raise AttributeError(f"module 'veadk.memory' has no attribute '{name}'")
+
+
+__all__ = ["ShortTermMemory", "LongTermMemory"]

veadk/memory/short_term_memory.py

@@ -15,8 +15,12 @@
 from functools import wraps
 from typing import Any, Callable, Literal
 
-from google.adk.sessions import DatabaseSessionService, InMemorySessionService
-from pydantic import BaseModel, Field
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+    InMemorySessionService,
+)
+from pydantic import BaseModel, Field, PrivateAttr
 
 from veadk.memory.short_term_memory_backends.mysql_backend import (
     MysqlSTMBackend,
@@ -61,10 +65,12 @@ class ShortTermMemory(BaseModel):
     after_load_memory_callback: Callable | None = None
     """A callback to be called after loading memory from the backend. The callback function should accept `Session` as an input."""
 
+    _session_service: BaseSessionService = PrivateAttr()
+
     def model_post_init(self, __context: Any) -> None:
         if self.db_url:
             logger.info("The `db_url` is set, ignore `backend` option.")
-            self.session_service = DatabaseSessionService(db_url=self.db_url)
+            self._session_service = DatabaseSessionService(db_url=self.db_url)
         else:
             if self.backend == "database":
                 logger.warning(
@@ -73,37 +79,41 @@ def model_post_init(self, __context: Any) -> None:
                 self.backend = "sqlite"
             match self.backend:
                 case "local":
-                    self.session_service = InMemorySessionService()
+                    self._session_service = InMemorySessionService()
                 case "mysql":
-                    self.session_service = MysqlSTMBackend(
+                    self._session_service = MysqlSTMBackend(
                         **self.backend_configs
                     ).session_service
                 case "sqlite":
-                    self.session_service = SQLiteSTMBackend(
+                    self._session_service = SQLiteSTMBackend(
                         local_path=self.local_database_path
                     ).session_service
                 case "redis":
-                    self.session_service = RedisSTMBackend(
+                    self._session_service = RedisSTMBackend(
                         **self.backend_configs
                     ).session_service
                 case "postgresql":
-                    self.session_service = PostgreSqlSTMBackend(
+                    self._session_service = PostgreSqlSTMBackend(
                         **self.backend_configs
                     ).session_service
 
         if self.after_load_memory_callback:
             wrap_get_session_with_callbacks(
-                self.session_service, self.after_load_memory_callback
+                self._session_service, self.after_load_memory_callback
             )
 
+    @property
+    def session_service(self) -> BaseSessionService:
+        return self._session_service
+
     async def create_session(
         self,
         app_name: str,
         user_id: str,
         session_id: str,
     ) -> None:
-        if isinstance(self.session_service, DatabaseSessionService):
-            list_sessions_response = await self.session_service.list_sessions(
+        if isinstance(self._session_service, DatabaseSessionService):
+            list_sessions_response = await self._session_service.list_sessions(
                 app_name=app_name, user_id=user_id
             )
 
@@ -112,12 +122,12 @@ async def create_session(
             )
 
         if (
-            await self.session_service.get_session(
+            await self._session_service.get_session(
                 app_name=app_name, user_id=user_id, session_id=session_id
             )
             is None
         ):
             # create a new session for this running
-            await self.session_service.create_session(
+            await self._session_service.create_session(
                 app_name=app_name, user_id=user_id, session_id=session_id
             )