veadk-python 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of veadk-python might be problematic. Click here for more details.
- veadk/agent.py +40 -8
- veadk/cli/cli_deploy.py +5 -1
- veadk/cli/cli_init.py +25 -6
- veadk/cloud/cloud_app.py +21 -6
- veadk/consts.py +33 -1
- veadk/database/database_adapter.py +88 -0
- veadk/database/kv/redis_database.py +47 -0
- veadk/database/local_database.py +22 -4
- veadk/database/relational/mysql_database.py +58 -0
- veadk/database/vector/opensearch_vector_database.py +6 -3
- veadk/database/viking/viking_database.py +72 -3
- veadk/integrations/ve_cr/__init__.py +13 -0
- veadk/integrations/ve_cr/ve_cr.py +205 -0
- veadk/integrations/ve_faas/template/cookiecutter.json +2 -1
- veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/clean.py +23 -0
- veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/src/app.py +28 -2
- veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/src/requirements.txt +3 -1
- veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/src/run.sh +5 -2
- veadk/integrations/ve_faas/ve_faas.py +2 -0
- veadk/integrations/ve_faas/web_template/cookiecutter.json +17 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/__init__.py +13 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/clean.py +23 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/config.yaml.example +2 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/deploy.py +41 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/Dockerfile +23 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/app.py +123 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/init_db.py +46 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/models.py +36 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/requirements.txt +4 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/run.sh +21 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/static/css/style.css +368 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/static/js/admin.js +0 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/admin/dashboard.html +21 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/admin/edit_post.html +24 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/admin/login.html +21 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/admin/posts.html +53 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/base.html +45 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/index.html +29 -0
- veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/post.html +14 -0
- veadk/integrations/ve_tos/ve_tos.py +238 -0
- veadk/knowledgebase/knowledgebase.py +8 -0
- veadk/runner.py +140 -34
- veadk/tools/builtin_tools/image_edit.py +236 -0
- veadk/tools/builtin_tools/image_generate.py +236 -0
- veadk/tools/builtin_tools/video_generate.py +326 -0
- veadk/tools/sandbox/browser_sandbox.py +19 -9
- veadk/tools/sandbox/code_sandbox.py +21 -11
- veadk/tools/sandbox/computer_sandbox.py +16 -9
- veadk/tracing/base_tracer.py +0 -19
- veadk/tracing/telemetry/attributes/extractors/common_attributes_extractors.py +5 -0
- veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py +311 -128
- veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py +20 -14
- veadk/tracing/telemetry/attributes/extractors/types.py +15 -4
- veadk/tracing/telemetry/exporters/inmemory_exporter.py +3 -0
- veadk/tracing/telemetry/opentelemetry_tracer.py +15 -6
- veadk/tracing/telemetry/telemetry.py +128 -24
- veadk/utils/misc.py +40 -0
- veadk/version.py +1 -1
- {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/METADATA +1 -1
- {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/RECORD +64 -37
- {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/WHEEL +0 -0
- {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/entry_points.txt +0 -0
- {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/licenses/LICENSE +0 -0
- {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/top_level.txt +0 -0
veadk/agent.py
CHANGED
|
@@ -28,9 +28,10 @@ from typing_extensions import Any
|
|
|
28
28
|
|
|
29
29
|
from veadk.config import getenv
|
|
30
30
|
from veadk.consts import (
|
|
31
|
-
DEFALUT_MODEL_AGENT_PROVIDER,
|
|
32
31
|
DEFAULT_MODEL_AGENT_API_BASE,
|
|
33
32
|
DEFAULT_MODEL_AGENT_NAME,
|
|
33
|
+
DEFAULT_MODEL_AGENT_PROVIDER,
|
|
34
|
+
DEFAULT_MODEL_EXTRA_CONFIG,
|
|
34
35
|
)
|
|
35
36
|
from veadk.evaluation import EvalSetRecorder
|
|
36
37
|
from veadk.knowledgebase import KnowledgeBase
|
|
@@ -64,7 +65,7 @@ class Agent(LlmAgent):
|
|
|
64
65
|
model_name: str = getenv("MODEL_AGENT_NAME", DEFAULT_MODEL_AGENT_NAME)
|
|
65
66
|
"""The name of the model for agent running."""
|
|
66
67
|
|
|
67
|
-
model_provider: str = getenv("MODEL_AGENT_PROVIDER",
|
|
68
|
+
model_provider: str = getenv("MODEL_AGENT_PROVIDER", DEFAULT_MODEL_AGENT_PROVIDER)
|
|
68
69
|
"""The provider of the model for agent running."""
|
|
69
70
|
|
|
70
71
|
model_api_base: str = getenv("MODEL_AGENT_API_BASE", DEFAULT_MODEL_AGENT_API_BASE)
|
|
@@ -73,6 +74,9 @@ class Agent(LlmAgent):
|
|
|
73
74
|
model_api_key: str = Field(default_factory=lambda: getenv("MODEL_AGENT_API_KEY"))
|
|
74
75
|
"""The api key of the model for agent running."""
|
|
75
76
|
|
|
77
|
+
model_extra_config: dict = Field(default_factory=dict)
|
|
78
|
+
"""The extra config to include in the model requests."""
|
|
79
|
+
|
|
76
80
|
tools: list[ToolUnion] = []
|
|
77
81
|
"""The tools provided to agent."""
|
|
78
82
|
|
|
@@ -96,11 +100,39 @@ class Agent(LlmAgent):
|
|
|
96
100
|
|
|
97
101
|
def model_post_init(self, __context: Any) -> None:
|
|
98
102
|
super().model_post_init(None) # for sub_agents init
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
103
|
+
|
|
104
|
+
# combine user model config with VeADK defaults
|
|
105
|
+
headers = DEFAULT_MODEL_EXTRA_CONFIG["extra_headers"].copy()
|
|
106
|
+
body = DEFAULT_MODEL_EXTRA_CONFIG["extra_body"].copy()
|
|
107
|
+
|
|
108
|
+
if self.model_extra_config:
|
|
109
|
+
user_headers = self.model_extra_config.get("extra_headers", {})
|
|
110
|
+
user_body = self.model_extra_config.get("extra_body", {})
|
|
111
|
+
|
|
112
|
+
headers |= user_headers
|
|
113
|
+
body |= user_body
|
|
114
|
+
|
|
115
|
+
self.model_extra_config |= {
|
|
116
|
+
"extra_headers": headers,
|
|
117
|
+
"extra_body": body,
|
|
118
|
+
}
|
|
119
|
+
|
|
120
|
+
logger.info(f"Model extra config: {self.model_extra_config}")
|
|
121
|
+
|
|
122
|
+
if not self.model:
|
|
123
|
+
self.model = LiteLlm(
|
|
124
|
+
model=f"{self.model_provider}/{self.model_name}",
|
|
125
|
+
api_key=self.model_api_key,
|
|
126
|
+
api_base=self.model_api_base,
|
|
127
|
+
**self.model_extra_config,
|
|
128
|
+
)
|
|
129
|
+
logger.debug(
|
|
130
|
+
f"LiteLLM client created with config: {self.model_extra_config}"
|
|
131
|
+
)
|
|
132
|
+
else:
|
|
133
|
+
logger.warning(
|
|
134
|
+
"You are trying to use your own LiteLLM client, some default request headers may be missing."
|
|
135
|
+
)
|
|
104
136
|
|
|
105
137
|
if self.knowledgebase:
|
|
106
138
|
from veadk.tools import load_knowledgebase_tool
|
|
@@ -117,7 +149,7 @@ class Agent(LlmAgent):
|
|
|
117
149
|
|
|
118
150
|
logger.info(f"{self.__class__.__name__} `{self.name}` init done.")
|
|
119
151
|
logger.debug(
|
|
120
|
-
f"Agent: {self.model_dump(include={'name', 'model_name', 'model_api_base', 'tools'
|
|
152
|
+
f"Agent: {self.model_dump(include={'name', 'model_name', 'model_api_base', 'tools'})}"
|
|
121
153
|
)
|
|
122
154
|
|
|
123
155
|
async def _run(
|
veadk/cli/cli_deploy.py
CHANGED
|
@@ -14,6 +14,7 @@
|
|
|
14
14
|
|
|
15
15
|
|
|
16
16
|
import click
|
|
17
|
+
from veadk.version import VERSION
|
|
17
18
|
|
|
18
19
|
TEMP_PATH = "/tmp"
|
|
19
20
|
|
|
@@ -29,7 +30,9 @@ TEMP_PATH = "/tmp"
|
|
|
29
30
|
default=None,
|
|
30
31
|
help="Volcengine secret key",
|
|
31
32
|
)
|
|
32
|
-
@click.option(
|
|
33
|
+
@click.option(
|
|
34
|
+
"--vefaas-app-name", required=True, help="Expected Volcengine FaaS application name"
|
|
35
|
+
)
|
|
33
36
|
@click.option(
|
|
34
37
|
"--veapig-instance-name", default="", help="Expected Volcengine APIG instance name"
|
|
35
38
|
)
|
|
@@ -92,6 +95,7 @@ def deploy(
|
|
|
92
95
|
"veapig_service_name": veapig_service_name,
|
|
93
96
|
"veapig_upstream_name": veapig_upstream_name,
|
|
94
97
|
"use_adk_web": use_adk_web,
|
|
98
|
+
"veadk_version": VERSION,
|
|
95
99
|
}
|
|
96
100
|
|
|
97
101
|
cookiecutter(
|
veadk/cli/cli_init.py
CHANGED
|
@@ -16,6 +16,8 @@ import warnings
|
|
|
16
16
|
from typing import Any
|
|
17
17
|
|
|
18
18
|
import click
|
|
19
|
+
from veadk.version import VERSION
|
|
20
|
+
|
|
19
21
|
|
|
20
22
|
warnings.filterwarnings(
|
|
21
23
|
"ignore", category=UserWarning, module="pydantic._internal._fields"
|
|
@@ -58,12 +60,21 @@ def _render_prompts() -> dict[str, Any]:
|
|
|
58
60
|
"veapig_service_name": veapig_service_name,
|
|
59
61
|
"veapig_upstream_name": veapig_upstream_name,
|
|
60
62
|
"use_adk_web": deploy_mode == "2",
|
|
63
|
+
"veadk_version": VERSION,
|
|
61
64
|
}
|
|
62
65
|
|
|
63
66
|
|
|
64
67
|
@click.command()
|
|
65
|
-
|
|
66
|
-
"""
|
|
68
|
+
@click.option(
|
|
69
|
+
"--vefaas-template-type", default="template", help="Expected template type"
|
|
70
|
+
)
|
|
71
|
+
def init(
|
|
72
|
+
vefaas_template_type: str,
|
|
73
|
+
) -> None:
|
|
74
|
+
"""Init a veadk project that can be deployed to Volcengine VeFaaS.
|
|
75
|
+
|
|
76
|
+
`template` is A2A/MCP/Web server template, `web_template` is for web applications (i.e., a simple blog).
|
|
77
|
+
"""
|
|
67
78
|
import shutil
|
|
68
79
|
from pathlib import Path
|
|
69
80
|
|
|
@@ -71,9 +82,14 @@ def init() -> None:
|
|
|
71
82
|
|
|
72
83
|
import veadk.integrations.ve_faas as vefaas
|
|
73
84
|
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
85
|
+
if vefaas_template_type == "web_template":
|
|
86
|
+
click.echo(
|
|
87
|
+
"Welcome use VeADK to create your project. We will generate a `simple-blog` web application for you."
|
|
88
|
+
)
|
|
89
|
+
else:
|
|
90
|
+
click.echo(
|
|
91
|
+
"Welcome use VeADK to create your project. We will generate a `weather-reporter` application for you."
|
|
92
|
+
)
|
|
77
93
|
|
|
78
94
|
cwd = Path.cwd()
|
|
79
95
|
local_dir_name = click.prompt("Local directory name", default="veadk-cloud-proj")
|
|
@@ -89,7 +105,10 @@ def init() -> None:
|
|
|
89
105
|
settings = _render_prompts()
|
|
90
106
|
settings["local_dir_name"] = local_dir_name
|
|
91
107
|
|
|
92
|
-
|
|
108
|
+
if not vefaas_template_type:
|
|
109
|
+
vefaas_template_type = "template"
|
|
110
|
+
|
|
111
|
+
template_dir_path = Path(vefaas.__file__).parent / vefaas_template_type
|
|
93
112
|
|
|
94
113
|
cookiecutter(
|
|
95
114
|
template=str(template_dir_path),
|
veadk/cloud/cloud_app.py
CHANGED
|
@@ -13,6 +13,7 @@
|
|
|
13
13
|
# limitations under the License.
|
|
14
14
|
|
|
15
15
|
import json
|
|
16
|
+
import time
|
|
16
17
|
from typing import Any
|
|
17
18
|
from uuid import uuid4
|
|
18
19
|
|
|
@@ -60,9 +61,11 @@ class CloudApp:
|
|
|
60
61
|
if not vefaas_endpoint:
|
|
61
62
|
self.vefaas_endpoint = self._get_vefaas_endpoint()
|
|
62
63
|
|
|
63
|
-
if
|
|
64
|
-
|
|
65
|
-
|
|
64
|
+
if (
|
|
65
|
+
self.vefaas_endpoint
|
|
66
|
+
and not self.vefaas_endpoint.startswith("http")
|
|
67
|
+
and not self.vefaas_endpoint.startswith("https")
|
|
68
|
+
):
|
|
66
69
|
raise ValueError(
|
|
67
70
|
f"Invalid endpoint: {vefaas_endpoint}. The endpoint must start with `http` or `https`."
|
|
68
71
|
)
|
|
@@ -92,12 +95,13 @@ class CloudApp:
|
|
|
92
95
|
raise ValueError(
|
|
93
96
|
f"VeFaaS CloudAPP with application_id `{self.vefaas_application_id}` or application_name `{self.vefaas_application_name}` not found."
|
|
94
97
|
)
|
|
95
|
-
cloud_resource = json.loads(app["CloudResource"])
|
|
96
98
|
|
|
97
99
|
try:
|
|
100
|
+
cloud_resource = json.loads(app["CloudResource"])
|
|
98
101
|
vefaas_endpoint = cloud_resource["framework"]["url"]["system_url"]
|
|
99
102
|
except Exception as e:
|
|
100
|
-
|
|
103
|
+
logger.warning(f"VeFaaS cloudAPP could not get endpoint. Error: {e}")
|
|
104
|
+
vefaas_endpoint = ""
|
|
101
105
|
return vefaas_endpoint
|
|
102
106
|
|
|
103
107
|
def _get_vefaas_application_id_by_name(self) -> str:
|
|
@@ -167,7 +171,18 @@ class CloudApp:
|
|
|
167
171
|
|
|
168
172
|
vefaas_client = VeFaaS(access_key=volcengine_ak, secret_key=volcengine_sk)
|
|
169
173
|
vefaas_client.delete(self.vefaas_application_id)
|
|
170
|
-
print(
|
|
174
|
+
print(
|
|
175
|
+
f"Cloud app {self.vefaas_application_id} delete request has been sent to VeFaaS"
|
|
176
|
+
)
|
|
177
|
+
while True:
|
|
178
|
+
try:
|
|
179
|
+
id = self._get_vefaas_application_id_by_name()
|
|
180
|
+
if not id:
|
|
181
|
+
break
|
|
182
|
+
time.sleep(3)
|
|
183
|
+
except Exception as _:
|
|
184
|
+
break
|
|
185
|
+
print("Delete application done.")
|
|
171
186
|
|
|
172
187
|
async def message_send(
|
|
173
188
|
self, message: str, session_id: str, user_id: str, timeout: float = 600.0
|
veadk/consts.py
CHANGED
|
@@ -12,6 +12,38 @@
|
|
|
12
12
|
# See the License for the specific language governing permissions and
|
|
13
13
|
# limitations under the License.
|
|
14
14
|
|
|
15
|
+
import time
|
|
16
|
+
|
|
17
|
+
from veadk.config import getenv
|
|
18
|
+
from veadk.version import VERSION
|
|
19
|
+
|
|
15
20
|
DEFAULT_MODEL_AGENT_NAME = "doubao-seed-1-6-250615"
|
|
16
|
-
|
|
21
|
+
DEFAULT_MODEL_AGENT_PROVIDER = "openai"
|
|
17
22
|
DEFAULT_MODEL_AGENT_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
|
|
23
|
+
DEFAULT_MODEL_EXTRA_CONFIG = {
|
|
24
|
+
"extra_headers": {
|
|
25
|
+
"x-is-encrypted": getenv("MODEL_AGENT_ENCRYPTED", "true"),
|
|
26
|
+
"veadk-source": "veadk",
|
|
27
|
+
"veadk-version": VERSION,
|
|
28
|
+
},
|
|
29
|
+
"extra_body": {
|
|
30
|
+
"caching": {
|
|
31
|
+
"type": getenv("MODEL_AGENT_CACHING", "enabled"),
|
|
32
|
+
},
|
|
33
|
+
"expire_at": int(time.time()) + 3600, # expire after 1 hour
|
|
34
|
+
},
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
DEFAULT_APMPLUS_OTEL_EXPORTER_ENDPOINT = "http://apmplus-cn-beijing.volces.com:4317"
|
|
38
|
+
DEFAULT_APMPLUS_OTEL_EXPORTER_SERVICE_NAME = "veadk_tracing"
|
|
39
|
+
|
|
40
|
+
DEFAULT_COZELOOP_OTEL_EXPORTER_ENDPOINT = (
|
|
41
|
+
"https://api.coze.cn/v1/loop/opentelemetry/v1/traces"
|
|
42
|
+
)
|
|
43
|
+
|
|
44
|
+
DEFAULT_TLS_OTEL_EXPORTER_ENDPOINT = "https://tls-cn-beijing.volces.com:4318/v1/traces"
|
|
45
|
+
DEFAULT_TLS_OTEL_EXPORTER_REGION = "cn-beijing"
|
|
46
|
+
|
|
47
|
+
DEFAULT_CR_INSTANCE_NAME = "veadk-user-instance"
|
|
48
|
+
DEFAULT_CR_NAMESPACE_NAME = "veadk-user-namespace"
|
|
49
|
+
DEFAULT_CR_REPO_NAME = "veadk-user-repo"
|
|
@@ -54,6 +54,33 @@ class KVDatabaseAdapter:
|
|
|
54
54
|
logger.error(f"Failed to search from Redis: index={index} error={e}")
|
|
55
55
|
raise e
|
|
56
56
|
|
|
57
|
+
def delete_doc(self, index: str, id: str) -> bool:
|
|
58
|
+
logger.debug(f"Deleting document from Redis database: index={index} id={id}")
|
|
59
|
+
try:
|
|
60
|
+
# For Redis, we need to handle deletion differently since RedisDatabase.delete_doc
|
|
61
|
+
# takes a key and a single id
|
|
62
|
+
result = self.client.delete_doc(key=index, id=id)
|
|
63
|
+
return result
|
|
64
|
+
except Exception as e:
|
|
65
|
+
logger.error(
|
|
66
|
+
f"Failed to delete document from Redis database: index={index} id={id} error={e}"
|
|
67
|
+
)
|
|
68
|
+
return False
|
|
69
|
+
|
|
70
|
+
def list_docs(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
|
|
71
|
+
logger.debug(f"Listing documents from Redis database: index={index}")
|
|
72
|
+
try:
|
|
73
|
+
# Get all documents from Redis
|
|
74
|
+
docs = self.client.list_docs(key=index)
|
|
75
|
+
|
|
76
|
+
# Apply offset and limit for pagination
|
|
77
|
+
return docs[offset : offset + limit]
|
|
78
|
+
except Exception as e:
|
|
79
|
+
logger.error(
|
|
80
|
+
f"Failed to list documents from Redis database: index={index} error={e}"
|
|
81
|
+
)
|
|
82
|
+
return []
|
|
83
|
+
|
|
57
84
|
|
|
58
85
|
class RelationalDatabaseAdapter:
|
|
59
86
|
def __init__(self, client):
|
|
@@ -108,6 +135,28 @@ class RelationalDatabaseAdapter:
|
|
|
108
135
|
|
|
109
136
|
return [item["data"] for item in results]
|
|
110
137
|
|
|
138
|
+
def delete_doc(self, index: str, id: str) -> bool:
|
|
139
|
+
logger.debug(f"Deleting document from SQL database: table_name={index} id={id}")
|
|
140
|
+
try:
|
|
141
|
+
# Convert single id to list for the client method
|
|
142
|
+
result = self.client.delete_doc(table=index, ids=[int(id)])
|
|
143
|
+
return result
|
|
144
|
+
except Exception as e:
|
|
145
|
+
logger.error(
|
|
146
|
+
f"Failed to delete document from SQL database: table_name={index} id={id} error={e}"
|
|
147
|
+
)
|
|
148
|
+
return False
|
|
149
|
+
|
|
150
|
+
def list_docs(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
|
|
151
|
+
logger.debug(f"Listing documents from SQL database: table_name={index}")
|
|
152
|
+
try:
|
|
153
|
+
return self.client.list_docs(table=index, offset=offset, limit=limit)
|
|
154
|
+
except Exception as e:
|
|
155
|
+
logger.error(
|
|
156
|
+
f"Failed to list documents from SQL database: table_name={index} error={e}"
|
|
157
|
+
)
|
|
158
|
+
return []
|
|
159
|
+
|
|
111
160
|
|
|
112
161
|
class VectorDatabaseAdapter:
|
|
113
162
|
def __init__(self, client):
|
|
@@ -152,6 +201,23 @@ class VectorDatabaseAdapter:
|
|
|
152
201
|
top_k=top_k,
|
|
153
202
|
)
|
|
154
203
|
|
|
204
|
+
def delete_doc(self, index: str, id: str) -> bool:
|
|
205
|
+
self._validate_index(index)
|
|
206
|
+
logger.debug(f"Deleting documents from vector database: index={index} id={id}")
|
|
207
|
+
try:
|
|
208
|
+
self.client.delete_by_id(collection_name=index, id=id)
|
|
209
|
+
return True
|
|
210
|
+
except Exception as e:
|
|
211
|
+
logger.error(
|
|
212
|
+
f"Failed to delete document from vector database: index={index} id={id} error={e}"
|
|
213
|
+
)
|
|
214
|
+
return False
|
|
215
|
+
|
|
216
|
+
def list_docs(self, index: str, offset: int = 0, limit: int = 1000) -> list[dict]:
|
|
217
|
+
self._validate_index(index)
|
|
218
|
+
logger.debug(f"Listing documents from vector database: index={index}")
|
|
219
|
+
return self.client.list_docs(collection_name=index, offset=offset, limit=limit)
|
|
220
|
+
|
|
155
221
|
|
|
156
222
|
class VikingDatabaseAdapter:
|
|
157
223
|
def __init__(self, client):
|
|
@@ -212,6 +278,16 @@ class VikingDatabaseAdapter:
|
|
|
212
278
|
|
|
213
279
|
return self.client.query(query, collection_name=index, top_k=top_k)
|
|
214
280
|
|
|
281
|
+
def delete_doc(self, index: str, id: str) -> bool:
|
|
282
|
+
self._validate_index(index)
|
|
283
|
+
logger.debug(f"Deleting documents from vector database: index={index} id={id}")
|
|
284
|
+
return self.client.delete_by_id(collection_name=index, id=id)
|
|
285
|
+
|
|
286
|
+
def list_docs(self, index: str, offset: int, limit: int) -> list[dict]:
|
|
287
|
+
self._validate_index(index)
|
|
288
|
+
logger.debug(f"Listing documents from vector database: index={index}")
|
|
289
|
+
return self.client.list_docs(collection_name=index, offset=offset, limit=limit)
|
|
290
|
+
|
|
215
291
|
|
|
216
292
|
class VikingMemoryDatabaseAdapter:
|
|
217
293
|
def __init__(self, client):
|
|
@@ -248,6 +324,12 @@ class VikingMemoryDatabaseAdapter:
|
|
|
248
324
|
result = self.client.query(query, collection_name=index, top_k=top_k, **kwargs)
|
|
249
325
|
return result
|
|
250
326
|
|
|
327
|
+
def delete_docs(self, index: str, ids: list[int]):
|
|
328
|
+
raise NotImplementedError("VikingMemoryDatabase does not support delete_docs")
|
|
329
|
+
|
|
330
|
+
def list_docs(self, index: str):
|
|
331
|
+
raise NotImplementedError("VikingMemoryDatabase does not support list_docs")
|
|
332
|
+
|
|
251
333
|
|
|
252
334
|
class LocalDatabaseAdapter:
|
|
253
335
|
def __init__(self, client):
|
|
@@ -261,6 +343,12 @@ class LocalDatabaseAdapter:
|
|
|
261
343
|
def query(self, query: str, **kwargs):
|
|
262
344
|
return self.client.query(query, **kwargs)
|
|
263
345
|
|
|
346
|
+
def delete_doc(self, index: str, id: str) -> bool:
|
|
347
|
+
return self.client.delete_doc(id)
|
|
348
|
+
|
|
349
|
+
def list_docs(self, index: str, offset: int = 0, limit: int = 100) -> list[dict]:
|
|
350
|
+
return self.client.list_docs(offset=offset, limit=limit)
|
|
351
|
+
|
|
264
352
|
|
|
265
353
|
MAPPING = {
|
|
266
354
|
"RedisDatabase": KVDatabaseAdapter,
|
|
@@ -110,3 +110,50 @@ class RedisDatabase(BaseModel, BaseDatabase):
|
|
|
110
110
|
except Exception as e:
|
|
111
111
|
logger.error(f"Failed to delete key `{key}`: {e}")
|
|
112
112
|
raise e
|
|
113
|
+
|
|
114
|
+
def delete_doc(self, key: str, id: str) -> bool:
|
|
115
|
+
"""Delete a specific document by ID from a Redis list.
|
|
116
|
+
|
|
117
|
+
Args:
|
|
118
|
+
key: The Redis key (list) to delete from
|
|
119
|
+
id: The ID of the document to delete
|
|
120
|
+
|
|
121
|
+
Returns:
|
|
122
|
+
bool: True if deletion was successful, False otherwise
|
|
123
|
+
"""
|
|
124
|
+
try:
|
|
125
|
+
# Get all items in the list
|
|
126
|
+
items = self._client.lrange(key, 0, -1)
|
|
127
|
+
|
|
128
|
+
# Find the index of the item to delete
|
|
129
|
+
for i, item in enumerate(items):
|
|
130
|
+
# Assuming the item is stored as a JSON string with an 'id' field
|
|
131
|
+
# If it's just the content, we'll use the list index as ID
|
|
132
|
+
if str(i) == id:
|
|
133
|
+
self._client.lrem(key, 1, item)
|
|
134
|
+
return True
|
|
135
|
+
|
|
136
|
+
logger.warning(f"Document with id {id} not found in key {key}")
|
|
137
|
+
return False
|
|
138
|
+
except Exception as e:
|
|
139
|
+
logger.error(f"Failed to delete document with id {id} from key {key}: {e}")
|
|
140
|
+
return False
|
|
141
|
+
|
|
142
|
+
def list_docs(self, key: str) -> list[dict]:
|
|
143
|
+
"""List all documents in a Redis list.
|
|
144
|
+
|
|
145
|
+
Args:
|
|
146
|
+
key: The Redis key (list) to list documents from
|
|
147
|
+
|
|
148
|
+
Returns:
|
|
149
|
+
list[dict]: List of documents with id and content
|
|
150
|
+
"""
|
|
151
|
+
try:
|
|
152
|
+
items = self._client.lrange(key, 0, -1)
|
|
153
|
+
return [
|
|
154
|
+
{"id": str(i), "content": item, "metadata": {}}
|
|
155
|
+
for i, item in enumerate(items)
|
|
156
|
+
]
|
|
157
|
+
except Exception as e:
|
|
158
|
+
logger.error(f"Failed to list documents from key {key}: {e}")
|
|
159
|
+
return []
|
veadk/database/local_database.py
CHANGED
|
@@ -24,20 +24,38 @@ class LocalDataBase(BaseDatabase):
|
|
|
24
24
|
|
|
25
25
|
def __init__(self, **kwargs):
|
|
26
26
|
super().__init__()
|
|
27
|
-
self.data =
|
|
27
|
+
self.data = {}
|
|
28
28
|
self._type = "local"
|
|
29
|
+
self._next_id = 0 # Used to generate unique IDs
|
|
29
30
|
|
|
30
31
|
def add_texts(self, texts: list[str], **kwargs):
|
|
31
|
-
|
|
32
|
+
for text in texts:
|
|
33
|
+
self.data[str(self._next_id)] = text
|
|
34
|
+
self._next_id += 1
|
|
32
35
|
|
|
33
36
|
def is_empty(self):
|
|
34
37
|
return len(self.data) == 0
|
|
35
38
|
|
|
36
39
|
def query(self, query: str, **kwargs: Any) -> list[str]:
|
|
37
|
-
return self.data
|
|
40
|
+
return list(self.data.values())
|
|
38
41
|
|
|
39
42
|
def delete(self, **kwargs: Any):
|
|
40
|
-
self.data =
|
|
43
|
+
self.data = {}
|
|
41
44
|
|
|
42
45
|
def add(self, texts: list[str], **kwargs: Any):
|
|
43
46
|
return self.add_texts(texts)
|
|
47
|
+
|
|
48
|
+
def list_docs(self, **kwargs: Any) -> list[dict]:
|
|
49
|
+
return [
|
|
50
|
+
{"id": id, "content": content, "metadata": {}}
|
|
51
|
+
for id, content in self.data.items()
|
|
52
|
+
]
|
|
53
|
+
|
|
54
|
+
def delete_doc(self, id: str, **kwargs: Any):
|
|
55
|
+
if id not in self.data:
|
|
56
|
+
raise ValueError(f"id {id} not found")
|
|
57
|
+
try:
|
|
58
|
+
del self.data[id]
|
|
59
|
+
return True
|
|
60
|
+
except Exception:
|
|
61
|
+
return False
|
|
@@ -111,5 +111,63 @@ class MysqlDatabase(BaseModel, BaseDatabase):
|
|
|
111
111
|
logger.error(f"Failed to drop table {table}: {e}")
|
|
112
112
|
raise e
|
|
113
113
|
|
|
114
|
+
def delete_doc(self, table: str, ids: list[int]) -> bool:
|
|
115
|
+
"""Delete documents by IDs from a MySQL table.
|
|
116
|
+
|
|
117
|
+
Args:
|
|
118
|
+
table: The table name to delete from
|
|
119
|
+
ids: List of document IDs to delete
|
|
120
|
+
|
|
121
|
+
Returns:
|
|
122
|
+
bool: True if deletion was successful, False otherwise
|
|
123
|
+
"""
|
|
124
|
+
if not self.table_exists(table):
|
|
125
|
+
logger.warning(f"Table {table} does not exist. Skipping delete operation.")
|
|
126
|
+
return False
|
|
127
|
+
|
|
128
|
+
if not ids:
|
|
129
|
+
return True # Nothing to delete
|
|
130
|
+
|
|
131
|
+
try:
|
|
132
|
+
with self._connection.cursor() as cursor:
|
|
133
|
+
# Create placeholders for the IDs
|
|
134
|
+
placeholders = ",".join(["%s"] * len(ids))
|
|
135
|
+
sql = f"DELETE FROM `{table}` WHERE id IN ({placeholders})"
|
|
136
|
+
cursor.execute(sql, ids)
|
|
137
|
+
self._connection.commit()
|
|
138
|
+
logger.info(f"Deleted {cursor.rowcount} documents from table {table}")
|
|
139
|
+
return True
|
|
140
|
+
except Exception as e:
|
|
141
|
+
logger.error(f"Failed to delete documents from table {table}: {e}")
|
|
142
|
+
return False
|
|
143
|
+
|
|
144
|
+
def list_docs(self, table: str, offset: int = 0, limit: int = 100) -> list[dict]:
|
|
145
|
+
"""List documents from a MySQL table.
|
|
146
|
+
|
|
147
|
+
Args:
|
|
148
|
+
table: The table name to list documents from
|
|
149
|
+
offset: Offset for pagination
|
|
150
|
+
limit: Limit for pagination
|
|
151
|
+
|
|
152
|
+
Returns:
|
|
153
|
+
list[dict]: List of documents with id and content
|
|
154
|
+
"""
|
|
155
|
+
if not self.table_exists(table):
|
|
156
|
+
logger.warning(f"Table {table} does not exist. Returning empty list.")
|
|
157
|
+
return []
|
|
158
|
+
|
|
159
|
+
try:
|
|
160
|
+
with self._connection.cursor() as cursor:
|
|
161
|
+
sql = f"SELECT id, data FROM `{table}` ORDER BY created_at DESC LIMIT %s OFFSET %s"
|
|
162
|
+
cursor.execute(sql, (limit, offset))
|
|
163
|
+
results = cursor.fetchall()
|
|
164
|
+
return [
|
|
165
|
+
{"id": str(row["id"]), "content": row["data"], "metadata": {}}
|
|
166
|
+
for row in results
|
|
167
|
+
]
|
|
168
|
+
except Exception as e:
|
|
169
|
+
logger.error(f"Failed to list documents from table {table}: {e}")
|
|
170
|
+
return []
|
|
171
|
+
|
|
114
172
|
def is_empty(self):
|
|
115
173
|
pass
|
|
@@ -219,7 +219,9 @@ class OpenSearchVectorDatabase(BaseModel, BaseDatabase):
|
|
|
219
219
|
response = self._opensearch_client.indices.get_alias()
|
|
220
220
|
return list(response.keys())
|
|
221
221
|
|
|
222
|
-
def
|
|
222
|
+
def list_docs(
|
|
223
|
+
self, collection_name: str, offset: int = 0, limit: int = 10000
|
|
224
|
+
) -> list[dict]:
|
|
223
225
|
"""Match all docs in one index of OpenSearch"""
|
|
224
226
|
if not self.collection_exists(collection_name):
|
|
225
227
|
logger.warning(
|
|
@@ -227,12 +229,13 @@ class OpenSearchVectorDatabase(BaseModel, BaseDatabase):
|
|
|
227
229
|
)
|
|
228
230
|
return []
|
|
229
231
|
|
|
230
|
-
query = {"size":
|
|
232
|
+
query = {"size": limit, "from": offset, "query": {"match_all": {}}}
|
|
231
233
|
response = self._opensearch_client.search(index=collection_name, body=query)
|
|
232
234
|
return [
|
|
233
235
|
{
|
|
234
236
|
"id": hit["_id"],
|
|
235
|
-
"
|
|
237
|
+
"content": hit["_source"]["page_content"],
|
|
238
|
+
"metadata": {},
|
|
236
239
|
}
|
|
237
240
|
for hit in response["hits"]["hits"]
|
|
238
241
|
]
|
|
@@ -41,6 +41,8 @@ get_collections_path = "/api/knowledge/collection/info"
|
|
|
41
41
|
doc_add_path = "/api/knowledge/doc/add"
|
|
42
42
|
doc_info_path = "/api/knowledge/doc/info"
|
|
43
43
|
doc_del_path = "/api/collection/drop"
|
|
44
|
+
list_docs_path = "/api/knowledge/point/list"
|
|
45
|
+
delete_docs_path = "/api/knowledge/point/delete"
|
|
44
46
|
|
|
45
47
|
|
|
46
48
|
class VolcengineTOSConfig(BaseModel):
|
|
@@ -387,9 +389,9 @@ class VikingDatabase(BaseModel, BaseDatabase):
|
|
|
387
389
|
logger.error(f"Error in list_collections: {result['message']}")
|
|
388
390
|
raise ValueError(f"Error in list_collections: {result['message']}")
|
|
389
391
|
|
|
390
|
-
collections = result["data"]
|
|
391
|
-
if
|
|
392
|
-
|
|
392
|
+
collections = result["data"].get("collection_list", [])
|
|
393
|
+
if len(collections) == 0:
|
|
394
|
+
return False
|
|
393
395
|
|
|
394
396
|
collection_list = set()
|
|
395
397
|
|
|
@@ -400,3 +402,70 @@ class VikingDatabase(BaseModel, BaseDatabase):
|
|
|
400
402
|
return True
|
|
401
403
|
else:
|
|
402
404
|
return False
|
|
405
|
+
|
|
406
|
+
def list_docs(
|
|
407
|
+
self, collection_name: str, offset: int = 0, limit: int = -1
|
|
408
|
+
) -> list[dict]:
|
|
409
|
+
request_params = {
|
|
410
|
+
"collection_name": collection_name,
|
|
411
|
+
"project": self.config.project,
|
|
412
|
+
"offset": offset,
|
|
413
|
+
"limit": limit,
|
|
414
|
+
}
|
|
415
|
+
|
|
416
|
+
create_collection_req = prepare_request(
|
|
417
|
+
method="POST",
|
|
418
|
+
path=list_docs_path,
|
|
419
|
+
config=self.config,
|
|
420
|
+
data=request_params,
|
|
421
|
+
)
|
|
422
|
+
resp = requests.request(
|
|
423
|
+
method=create_collection_req.method,
|
|
424
|
+
url="https://{}{}".format(
|
|
425
|
+
g_knowledge_base_domain, create_collection_req.path
|
|
426
|
+
),
|
|
427
|
+
headers=create_collection_req.headers,
|
|
428
|
+
data=create_collection_req.body,
|
|
429
|
+
)
|
|
430
|
+
|
|
431
|
+
result = resp.json()
|
|
432
|
+
if result["code"] != 0:
|
|
433
|
+
logger.error(f"Error in list_docs: {result['message']}")
|
|
434
|
+
raise ValueError(f"Error in list_docs: {result['message']}")
|
|
435
|
+
|
|
436
|
+
data = [
|
|
437
|
+
{
|
|
438
|
+
"id": res["point_id"],
|
|
439
|
+
"content": res["content"],
|
|
440
|
+
"metadata": res["doc_info"],
|
|
441
|
+
}
|
|
442
|
+
for res in result["data"]["point_list"]
|
|
443
|
+
]
|
|
444
|
+
return data
|
|
445
|
+
|
|
446
|
+
def delete_by_id(self, collection_name: str, id: str) -> bool:
|
|
447
|
+
request_params = {
|
|
448
|
+
"collection_name": collection_name,
|
|
449
|
+
"project": self.config.project,
|
|
450
|
+
"point_id": id,
|
|
451
|
+
}
|
|
452
|
+
|
|
453
|
+
create_collection_req = prepare_request(
|
|
454
|
+
method="POST",
|
|
455
|
+
path=delete_docs_path,
|
|
456
|
+
config=self.config,
|
|
457
|
+
data=request_params,
|
|
458
|
+
)
|
|
459
|
+
resp = requests.request(
|
|
460
|
+
method=create_collection_req.method,
|
|
461
|
+
url="https://{}{}".format(
|
|
462
|
+
g_knowledge_base_domain, create_collection_req.path
|
|
463
|
+
),
|
|
464
|
+
headers=create_collection_req.headers,
|
|
465
|
+
data=create_collection_req.body,
|
|
466
|
+
)
|
|
467
|
+
|
|
468
|
+
result = resp.json()
|
|
469
|
+
if result["code"] != 0:
|
|
470
|
+
return False
|
|
471
|
+
return True
|