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

veadk/database/vector/opensearch_vector_database.py

@@ -186,7 +186,7 @@ class OpenSearchVectorDatabase(BaseModel, BaseDatabase):
         assert collection_name is not None, "Collection name is required."
         if not self._opensearch_client.indices.exists(index=collection_name):
             logger.warning(
-                f"querying {query}, but collection {collection_name} does not exist. retun a empty list."
+                f"querying {query}, but collection {collection_name} does not exist. return a empty list."
             )
             return []
         query_vector = self._embedding_client.embed_query(query)
@@ -196,10 +196,58 @@ class OpenSearchVectorDatabase(BaseModel, BaseDatabase):
 
     @override
     def delete(self, collection_name: str, **kwargs: Any):
+        """drop index"""
         if not self._opensearch_client.indices.exists(index=collection_name):
             raise ValueError(f"Collection {collection_name} does not exist.")
         self._opensearch_client.indices.delete(index=collection_name)
 
-    def is_empty(self, collection_name: str):
+    def is_empty(self, collection_name: str) -> bool:
         response = self._opensearch_client.count(index=collection_name)
         return response["count"] == 0
+
+    def collection_exists(self, collection_name: str) -> bool:
+        return self._opensearch_client.indices.exists(index=collection_name)
+
+    def list_all_collection(self) -> list:
+        """List all index name of OpenSearch."""
+        response = self._opensearch_client.indices.get_alias()
+        return list(response.keys())
+
+    def get_all_docs(self, collection_name: str, size: int = 10000) -> list[dict]:
+        """Match all docs in one index of OpenSearch"""
+        if not self.collection_exists(collection_name):
+            logger.warning(
+                f"Get all docs, but collection {collection_name} does not exist. return a empty list."
+            )
+            return []
+
+        query = {"size": size, "query": {"match_all": {}}}
+        response = self._opensearch_client.search(index=collection_name, body=query)
+        return [
+            {
+                "id": hit["_id"],
+                "page_content": hit["_source"]["page_content"],
+            }
+            for hit in response["hits"]["hits"]
+        ]
+
+    def delete_by_query(self, collection_name: str, query: str):
+        """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}}}
+        response = self._opensearch_client.delete_by_query(
+            index=collection_name, body=query
+        )
+        self._opensearch_client.indices.refresh(index=collection_name)
+        return response
+
+    def delete_by_id(self, collection_name: str, id: str):
+        """Delete docs by id in index of OpenSearch"""
+        if not self.collection_exists(collection_name):
+            raise ValueError(f"Collection {collection_name} does not exist.")
+
+        response = self._opensearch_client.delete(index=collection_name, id=id)
+        self._opensearch_client.indices.refresh(index=collection_name)
+        return response

veadk/database/viking/viking_database.py

@@ -40,6 +40,7 @@ list_collections_path = "/api/knowledge/collection/list"
 get_collections_path = "/api/knowledge/collection/info"
 doc_add_path = "/api/knowledge/doc/add"
 doc_info_path = "/api/knowledge/doc/info"
+doc_del_path = "/api/collection/drop"
 
 
 class VolcengineTOSConfig(BaseModel):
@@ -215,7 +216,12 @@ class VikingDatabase(BaseModel, BaseDatabase):
 
         return doc_id
 
-    def add(self, data: str | list[str] | TextIO | BinaryIO | bytes, **kwargs: Any):
+    def add(
+        self,
+        data: str | list[str] | TextIO | BinaryIO | bytes,
+        collection_name: str,
+        **kwargs,
+    ):
         """
         Args:
             data: str, file path or file stream:  Both file or file.read() are acceptable.
@@ -226,8 +232,6 @@ class VikingDatabase(BaseModel, BaseDatabase):
                 "doc_id": "<doc_id>",
             }
         """
-        collection_name = kwargs.get("collection_name")
-        assert collection_name is not None, "collection_name is required"
 
         status, tos_url = self._upload_to_tos(data=data, **kwargs)
         if status != 200:
@@ -243,9 +247,23 @@ class VikingDatabase(BaseModel, BaseDatabase):
         }
 
     def delete(self, **kwargs: Any):
-        # collection_name = kwargs.get("collection_name")
-        # todo: delete vikingdb
-        ...
+        collection_name = kwargs.get("collection_name")
+        resource_id = kwargs.get("resource_id")
+        request_param = {"collection_name": collection_name, "resource_id": resource_id}
+        doc_del_req = prepare_request(
+            method="POST", path=doc_del_path, config=self.config, data=request_param
+        )
+        rsp = requests.request(
+            method=doc_del_req.method,
+            url="http://{}{}".format(g_knowledge_base_domain, doc_del_req.path),
+            headers=doc_del_req.headers,
+            data=doc_del_req.body,
+        )
+        result = rsp.json()
+        if result["code"] != 0:
+            logger.error(f"Error in add_doc: {result['message']}")
+            return {"error": result["message"]}
+        return {}
 
     def query(self, query: str, **kwargs: Any) -> list[str]:
         """

veadk/database/viking/viking_memory_db.py

@@ -53,7 +53,8 @@ class VikingMemConfig(BaseModel):
     )
 
 
-class VikingDBMemoryException(Exception):
+# ======= adapted from https://github.com/volcengine/mcp-server/blob/main/server/mcp_server_vikingdb_memory/src/mcp_server_vikingdb_memory/common/memory_client.py =======
+class VikingMemoryException(Exception):
     def __init__(self, code, request_id, message=None):
         self.code = code
         self.request_id = request_id
@@ -65,15 +66,15 @@ class VikingDBMemoryException(Exception):
         return self.message
 
 
-class VikingDBMemoryService(Service):
+class VikingMemoryService(Service):
     _instance_lock = threading.Lock()
 
     def __new__(cls, *args, **kwargs):
-        if not hasattr(VikingDBMemoryService, "_instance"):
-            with VikingDBMemoryService._instance_lock:
-                if not hasattr(VikingDBMemoryService, "_instance"):
-                    VikingDBMemoryService._instance = object.__new__(cls)
-        return VikingDBMemoryService._instance
+        if not hasattr(VikingMemoryService, "_instance"):
+            with VikingMemoryService._instance_lock:
+                if not hasattr(VikingMemoryService, "_instance"):
+                    VikingMemoryService._instance = object.__new__(cls)
+        return VikingMemoryService._instance
 
     def __init__(
         self,
@@ -86,11 +87,11 @@ class VikingDBMemoryService(Service):
         connection_timeout=30,
         socket_timeout=30,
     ):
-        self.service_info = VikingDBMemoryService.get_service_info(
+        self.service_info = VikingMemoryService.get_service_info(
             host, region, scheme, connection_timeout, socket_timeout
         )
-        self.api_info = VikingDBMemoryService.get_api_info()
-        super(VikingDBMemoryService, self).__init__(self.service_info, self.api_info)
+        self.api_info = VikingMemoryService.get_api_info()
+        super(VikingMemoryService, self).__init__(self.service_info, self.api_info)
         if ak:
             self.set_ak(ak)
         if sk:
@@ -100,12 +101,12 @@ class VikingDBMemoryService(Service):
         try:
             self.get_body("Ping", {}, json.dumps({}))
         except Exception as e:
-            raise VikingDBMemoryException(
+            raise VikingMemoryException(
                 1000028, "missed", "host or region is incorrect: {}".format(str(e))
             ) from None
 
     def setHeader(self, header):
-        api_info = VikingDBMemoryService.get_api_info()
+        api_info = VikingMemoryService.get_api_info()
         for key in api_info:
             for item in header:
                 api_info[key].header[item] = header[item]
@@ -211,17 +212,17 @@ class VikingDBMemoryService(Service):
             try:
                 res_json = json.loads(e.args[0].decode("utf-8"))
             except Exception:
-                raise VikingDBMemoryException(
+                raise VikingMemoryException(
                     1000028, "missed", "json load res error, res:{}".format(str(e))
                 ) from None
             code = res_json.get("code", 1000028)
             request_id = res_json.get("request_id", 1000028)
             message = res_json.get("message", None)
 
-            raise VikingDBMemoryException(code, request_id, message)
+            raise VikingMemoryException(code, request_id, message)
 
         if res == "":
-            raise VikingDBMemoryException(
+            raise VikingMemoryException(
                 1000028,
                 "missed",
                 "empty response due to unknown error, please contact customer service",
@@ -235,15 +236,15 @@ class VikingDBMemoryService(Service):
             try:
                 res_json = json.loads(e.args[0].decode("utf-8"))
             except Exception:
-                raise VikingDBMemoryException(
+                raise VikingMemoryException(
                     1000028, "missed", "json load res error, res:{}".format(str(e))
                 ) from None
             code = res_json.get("code", 1000028)
             request_id = res_json.get("request_id", 1000028)
             message = res_json.get("message", None)
-            raise VikingDBMemoryException(code, request_id, message)
+            raise VikingMemoryException(code, request_id, message)
         if res == "":
-            raise VikingDBMemoryException(
+            raise VikingMemoryException(
                 1000028,
                 "missed",
                 "empty response due to unknown error, please contact customer service",
@@ -363,14 +364,17 @@ def format_milliseconds(timestamp_ms):
     return dt.strftime("%Y%m%d %H:%M:%S")
 
 
-class VikingDatabaseMemory(BaseModel, BaseDatabase):
+# ======= adapted from https://github.com/volcengine/mcp-server/blob/main/server/mcp_server_vikingdb_memory/src/mcp_server_vikingdb_memory/common/memory_client.py =======
+
+
+class VikingMemoryDatabase(BaseModel, BaseDatabase):
     config: VikingMemConfig = Field(
         default_factory=VikingMemConfig,
         description="VikingDB configuration",
     )
 
     def model_post_init(self, context: Any, /) -> None:
-        self._vm = VikingDBMemoryService(
+        self._vm = VikingMemoryService(
             ak=self.config.volcengine_ak, sk=self.config.volcengine_sk
         )
 
@@ -511,8 +515,8 @@ class VikingDatabaseMemory(BaseModel, BaseDatabase):
         assert collection_name is not None, "collection_name is required"
         user_id = kwargs.get("user_id")
         assert user_id is not None, "user_id is required"
-
-        resp = self.search_memory(collection_name, query, user_id=user_id)
+        top_k = kwargs.get("top_k", 5)
+        resp = self.search_memory(collection_name, query, user_id=user_id, top_k=top_k)
         return resp
 
     def delete(self, **kwargs: Any):

veadk/evaluation/base_evaluator.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 
+import json
 import time
 import uuid
 from abc import abstractmethod
@@ -24,6 +25,8 @@ from google.adk.sessions import InMemorySessionService
 from google.genai import types
 from pydantic import BaseModel
 
+from veadk.utils.misc import formatted_timestamp
+
 
 class InvocationTestData(BaseModel):
     invocation_id: str = ""
@@ -79,15 +82,139 @@ class BaseEvaluator:
         self.result_list: list[EvalResultData] = []
         self.agent_information_list: list[dict] = []
 
-    def load_eval_set(self, eval_set_file: str) -> list[EvalSet]:
+    def _load_eval_set(self, eval_set_file: str) -> EvalSet:
         from .eval_set_file_loader import load_eval_set_from_file
 
         return load_eval_set_from_file(eval_set_file)
 
-    def generate_eval_data(self, eval_set_file_path: str):
+    def _load_eval_set_from_tracing(self, tracing_file: str) -> EvalSet:
+        try:
+            with open(tracing_file, "r") as f:
+                tracing_data = json.load(f)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON format in file {tracing_file}: {e}")
+        except Exception as e:
+            raise ValueError(f"Error reading file {tracing_file}: {e}")
+
+        # Group spans by trace_id
+        trace_groups = {}
+        for span in tracing_data:
+            trace_id = span["trace_id"]
+            if trace_id not in trace_groups:
+                trace_groups[trace_id] = []
+            trace_groups[trace_id].append(span)
+
+        # Convert to evalset format
+        eval_cases, conversation = [], []
+        app_name, user_id = "", ""
+        creation_timestamp = 0
+        for trace_id, spans in trace_groups.items():
+            tool_uses = []
+
+            # Extract tool_uses from spans with name starting with "execute_tool"
+            for span in spans:
+                if span["name"].startswith("execute_tool"):
+                    tool_uses.append(
+                        {
+                            "id": span["attributes"].get("gen_ai.tool.call.id", None),
+                            "args": json.loads(
+                                span["attributes"].get(
+                                    "gcp.vertex.agent.tool_call_args", "{}"
+                                )
+                            ),
+                            "name": span["attributes"].get("gen_ai.tool.name", None),
+                        }
+                    )
+
+            # Extract conversation data from spans with name starting with "invocation"
+            for span in spans:
+                if span["name"].startswith("invocation"):
+                    # Parse input.value and output.value as JSON
+                    input_value = json.loads(
+                        span["attributes"].get("input.value", "{}")
+                    )
+                    output_value = json.loads(
+                        span["attributes"].get("output.value", "{}")
+                    )
+
+                    user_content = json.loads(input_value.get("new_message", {}))
+                    final_response = json.loads(json.dumps(user_content))
+                    final_response["parts"][0]["text"] = (
+                        output_value.get("content", {})
+                        .get("parts", [{}])[0]
+                        .get("text", None)
+                    )
+                    final_response["role"] = None
+                    conversation.append(
+                        {
+                            "invocation_id": output_value.get(
+                                "invocation_id", str(uuid.uuid4())
+                            ),
+                            "user_content": user_content,
+                            "final_response": final_response,
+                            "intermediate_data": {
+                                "tool_uses": tool_uses,
+                                "intermediate_responses": [],
+                            },
+                            "creation_timestamp": span["start_time"] / 1e9,
+                        }
+                    )
+                    user_id = input_value.get("user_id", None)
+                    app_name = (
+                        span["name"].replace("invocation", "").strip().strip("[]")
+                    )
+                    creation_timestamp = span["start_time"] / 1e9
+
+        eval_cases.append(
+            {
+                "eval_id": f"veadk_eval_{formatted_timestamp()}",
+                "conversation": conversation,
+                "session_input": {
+                    "app_name": app_name,
+                    "user_id": user_id,
+                    "state": {},
+                },
+                "creation_timestamp": creation_timestamp,
+            }
+        )
+
+        evalset = EvalSet(
+            eval_set_id="default",
+            name="default",
+            description=None,
+            eval_cases=eval_cases,
+            creation_timestamp=creation_timestamp,
+        )
+
+        return evalset
+
+    def generate_eval_data(self, file_path: str):
+        """Generate evaluation data from a given file and assign it to the class attribute `invocation_list`."""
         eval_case_data_list: list[EvalCaseData] = []
 
-        eval_cases = self.load_eval_set(eval_set_file_path).eval_cases
+        try:
+            with open(file_path, "r") as f:
+                file_content = json.load(f)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON format in file {file_path}: {e}")
+        except Exception as e:
+            raise ValueError(f"Error reading file {file_path}: {e}")
+
+        if isinstance(file_content, dict) and "eval_cases" in file_content:
+            eval_cases = self._load_eval_set(file_path).eval_cases
+        elif (
+            isinstance(file_content, list)
+            and len(file_content) > 0
+            and all(
+                isinstance(span, dict) and "trace_id" in span for span in file_content
+            )
+        ):
+            eval_cases = self._load_eval_set_from_tracing(file_path).eval_cases
+        else:
+            raise ValueError(
+                f"Unsupported file format in {file_path}. Please provide a valid file."
+            )
+
         for eval_case in eval_cases:
             eval_case_data = EvalCaseData(invocations=[])
             self.agent_information_list.append(

veadk/evaluation/deepeval_evaluator/deepeval_evaluator.py

@@ -15,8 +15,8 @@
 import time
 
 from deepeval import evaluate
+from deepeval.evaluate import CacheConfig
 from deepeval.evaluate.types import EvaluationResult
-from deepeval.key_handler import KEY_FILE_HANDLER, KeyValues
 from deepeval.metrics import BaseMetric
 from deepeval.models import LocalModel
 from deepeval.test_case import LLMTestCase
@@ -24,15 +24,11 @@ from deepeval.test_case.llm_test_case import ToolCall
 from typing_extensions import override
 
 from veadk.config import getenv
+from veadk.evaluation.types import EvalResultCaseData, EvalResultMetadata
 from veadk.utils.logger import get_logger
 
 from ..base_evaluator import BaseEvaluator, EvalResultData, MetricResult
-from ..utils.prometheus import (
-    EvalResultCaseData,
-    EvalResultMetadata,
-    PrometheusPushgatewayConfig,
-    push_to_prometheus,
-)
+from ..utils.prometheus import PrometheusPushgatewayConfig, push_to_prometheus
 
 logger = get_logger(__name__)
 
@@ -56,32 +52,19 @@ class DeepevalEvaluator(BaseEvaluator):
             "https://ark.cn-beijing.volces.com/api/v3/",
         ),
         name: str = "veadk_deepeval_evaluator",
-        prometheus_config: PrometheusPushgatewayConfig = None,
+        prometheus_config: PrometheusPushgatewayConfig | None = None,
     ):
         super().__init__(agent=agent, name=name)
 
         self.judge_model_name = judge_model_name
-        self.judge_model = self.create_judge_model(
-            model_name=judge_model_name,
+        self.judge_model = LocalModel(
+            model=judge_model_name,
+            base_url=judge_model_api_base,
             api_key=judge_model_api_key,
-            api_base=judge_model_api_base,
         )
 
         self.prometheus_config = prometheus_config
 
-    def create_judge_model(
-        self,
-        model_name: str,
-        api_key: str,
-        api_base: str,
-    ):
-        KEY_FILE_HANDLER.write_key(KeyValues.LOCAL_MODEL_NAME, model_name)
-        KEY_FILE_HANDLER.write_key(KeyValues.LOCAL_MODEL_BASE_URL, api_base)
-        KEY_FILE_HANDLER.write_key(KeyValues.LOCAL_MODEL_API_KEY, api_key)
-        KEY_FILE_HANDLER.write_key(KeyValues.USE_LOCAL_MODEL, "YES")
-        KEY_FILE_HANDLER.write_key(KeyValues.USE_AZURE_OPENAI, "NO")
-        return LocalModel()
-
     @override
     async def eval(
         self,
@@ -90,9 +73,9 @@ class DeepevalEvaluator(BaseEvaluator):
         eval_id: str = f"test_{formatted_timestamp()}",
     ):
         """Target to Google ADK, we will use the same evaluation case format as Google ADK."""
-
         # Get evaluation data by parsing eval set file
         self.generate_eval_data(eval_set_file_path)
+
         # Get actual data by running agent
         logger.info("Start to run agent for actual data.")
         await self._run_agent_for_actual_data()
@@ -141,7 +124,11 @@ class DeepevalEvaluator(BaseEvaluator):
 
         # Run Deepeval evaluation according to metrics
         logger.info("Start to run Deepeval evaluation according to metrics.")
-        test_results = evaluate(test_cases=test_cases, metrics=metrics)
+        test_results = evaluate(
+            test_cases=test_cases,
+            metrics=metrics,
+            cache_config=CacheConfig(write_cache=False),
+        )
         for test_result in test_results.test_results:
             eval_result_data = EvalResultData(metric_results=[])
             for metrics_data_item in test_result.metrics_data:

veadk/evaluation/types.py

@@ -0,0 +1,33 @@
+# 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 dataclasses import dataclass
+
+
+@dataclass
+class EvalResultCaseData:
+    id: str
+    input: str
+    actual_output: str
+    expected_output: str
+    score: str
+    reason: str
+    status: str  # `PASSED` or `FAILURE`
+    latency: str
+
+
+@dataclass
+class EvalResultMetadata:
+    tested_model: str
+    judge_model: str

veadk/evaluation/utils/prometheus.py

@@ -12,38 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
 
 from prometheus_client import CollectorRegistry, Gauge, push_to_gateway
 from prometheus_client.exposition import basic_auth_handler
+from pydantic import Field
 
 from veadk.config import getenv
-
-
-@dataclass
-class EvalResultCaseData:
-    id: str
-    input: str
-    actual_output: str
-    expected_output: str
-    score: str
-    reason: str
-    status: str  # `PASSED` or `FAILURE`
-    latency: str
-
-
-@dataclass
-class EvalResultMetadata:
-    tested_model: str
-    judge_model: str
+from veadk.evaluation.types import EvalResultCaseData, EvalResultMetadata
 
 
 class PrometheusPushgatewayConfig:
-    url: str = getenv(
-        "OBSERVABILITY_PROMETHEUS_PUSHGATEWAY_URL",
+    url: str = Field(
+        ...,
+        default_factory=lambda: getenv(
+            "OBSERVABILITY_PROMETHEUS_PUSHGATEWAY_URL",
+        ),
+    )
+    username: str = Field(
+        ...,
+        default_factory=lambda: getenv(
+            "OBSERVABILITY_PROMETHEUS_USERNAME",
+        ),
+    )
+    password: str = Field(
+        ...,
+        default_factory=lambda: getenv(
+            "OBSERVABILITY_PROMETHEUS_PASSWORD",
+        ),
     )
-    username: str = getenv("OBSERVABILITY_PROMETHEUS_USERNAME")
-    password: str = getenv("OBSERVABILITY_PROMETHEUS_PASSWORD")
 
 
 registry = CollectorRegistry()