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.

Files changed (64) hide show
  1. veadk/agent.py +40 -8
  2. veadk/cli/cli_deploy.py +5 -1
  3. veadk/cli/cli_init.py +25 -6
  4. veadk/cloud/cloud_app.py +21 -6
  5. veadk/consts.py +33 -1
  6. veadk/database/database_adapter.py +88 -0
  7. veadk/database/kv/redis_database.py +47 -0
  8. veadk/database/local_database.py +22 -4
  9. veadk/database/relational/mysql_database.py +58 -0
  10. veadk/database/vector/opensearch_vector_database.py +6 -3
  11. veadk/database/viking/viking_database.py +72 -3
  12. veadk/integrations/ve_cr/__init__.py +13 -0
  13. veadk/integrations/ve_cr/ve_cr.py +205 -0
  14. veadk/integrations/ve_faas/template/cookiecutter.json +2 -1
  15. veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/clean.py +23 -0
  16. veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/src/app.py +28 -2
  17. veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/src/requirements.txt +3 -1
  18. veadk/integrations/ve_faas/template/{{cookiecutter.local_dir_name}}/src/run.sh +5 -2
  19. veadk/integrations/ve_faas/ve_faas.py +2 -0
  20. veadk/integrations/ve_faas/web_template/cookiecutter.json +17 -0
  21. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/__init__.py +13 -0
  22. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/clean.py +23 -0
  23. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/config.yaml.example +2 -0
  24. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/deploy.py +41 -0
  25. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/Dockerfile +23 -0
  26. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/app.py +123 -0
  27. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/init_db.py +46 -0
  28. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/models.py +36 -0
  29. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/requirements.txt +4 -0
  30. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/run.sh +21 -0
  31. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/static/css/style.css +368 -0
  32. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/static/js/admin.js +0 -0
  33. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/admin/dashboard.html +21 -0
  34. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/admin/edit_post.html +24 -0
  35. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/admin/login.html +21 -0
  36. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/admin/posts.html +53 -0
  37. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/base.html +45 -0
  38. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/index.html +29 -0
  39. veadk/integrations/ve_faas/web_template/{{cookiecutter.local_dir_name}}/src/templates/post.html +14 -0
  40. veadk/integrations/ve_tos/ve_tos.py +238 -0
  41. veadk/knowledgebase/knowledgebase.py +8 -0
  42. veadk/runner.py +140 -34
  43. veadk/tools/builtin_tools/image_edit.py +236 -0
  44. veadk/tools/builtin_tools/image_generate.py +236 -0
  45. veadk/tools/builtin_tools/video_generate.py +326 -0
  46. veadk/tools/sandbox/browser_sandbox.py +19 -9
  47. veadk/tools/sandbox/code_sandbox.py +21 -11
  48. veadk/tools/sandbox/computer_sandbox.py +16 -9
  49. veadk/tracing/base_tracer.py +0 -19
  50. veadk/tracing/telemetry/attributes/extractors/common_attributes_extractors.py +5 -0
  51. veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py +311 -128
  52. veadk/tracing/telemetry/attributes/extractors/tool_attributes_extractors.py +20 -14
  53. veadk/tracing/telemetry/attributes/extractors/types.py +15 -4
  54. veadk/tracing/telemetry/exporters/inmemory_exporter.py +3 -0
  55. veadk/tracing/telemetry/opentelemetry_tracer.py +15 -6
  56. veadk/tracing/telemetry/telemetry.py +128 -24
  57. veadk/utils/misc.py +40 -0
  58. veadk/version.py +1 -1
  59. {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/METADATA +1 -1
  60. {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/RECORD +64 -37
  61. {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/WHEEL +0 -0
  62. {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/entry_points.txt +0 -0
  63. {veadk_python-0.2.4.dist-info → veadk_python-0.2.6.dist-info}/licenses/LICENSE +0 -0
  64. {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", DEFALUT_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
- self.model = LiteLlm(
100
- model=f"{self.model_provider}/{self.model_name}",
101
- api_key=self.model_api_key,
102
- api_base=self.model_api_base,
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', 'serve_url'})}"
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("--vefaas-app-name", help="Expected Volcengine FaaS application name")
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
- def init() -> None:
66
- """Init a veadk project that can be deployed to Volcengine VeFaaS."""
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
- click.echo(
75
- "Welcome use VeADK to create your project. We will generate a `weather-reporter` application for you."
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
- template_dir_path = Path(vefaas.__file__).parent / "template"
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 not self.vefaas_endpoint.startswith(
64
- "http"
65
- ) and not self.vefaas_endpoint.startswith("https"):
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
- raise ValueError(f"VeFaaS cloudAPP could not get endpoint. Error: {e}")
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(f"Cloud app {self.vefaas_application_id} is deleting...")
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
- DEFALUT_MODEL_AGENT_PROVIDER = "openai"
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 []
@@ -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
- self.data.extend(texts)
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 get_all_docs(self, collection_name: str, size: int = 10000) -> list[dict]:
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": size, "query": {"match_all": {}}}
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
- "page_content": hit["_source"]["page_content"],
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"]["collection_list"]
391
- if not collections:
392
- raise ValueError(f"No collections found in project {self.config.project}.")
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