api(kb): use id instead of kb_id

pingcap · Mar 3, 2025 · 3aac55b · 3aac55b
1 parent c44a109
commit 3aac55b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/core/autoflow/main.py b/core/autoflow/main.py
@@ -42,14 +42,14 @@ def create_knowledge_base(
         embedding_model: EmbeddingModel,
         description: Optional[str] = None,
         index_methods: Optional[List[IndexMethod]] = None,
-        kb_id: Optional[uuid.UUID] = None,
+        id: Optional[uuid.UUID] = None,
     ) -> KnowledgeBase:
         return KnowledgeBase(
             name=name,
             description=description,
             index_methods=index_methods,
             chat_model=chat_model,
             embedding_model=embedding_model,
-            id=kb_id,
+            id=id,
             db_engine=self._db_engine,
         )
diff --git a/core/tests/test_knowledge_base.py b/core/tests/test_knowledge_base.py
@@ -38,7 +38,7 @@ def test_create_knowledge_base():
 
 def test_import_documents_from_files():
     kb = af.create_knowledge_base(
-        kb_id=uuid.UUID("01973588-65aa-4954-99fd-71eb5ecce167"),
+        id=uuid.UUID("01973588-65aa-4954-99fd-71eb5ecce167"),
         name="Test",
         description="This is a knowledge base for testing",
         index_methods=[IndexMethod.VECTOR_SEARCH, IndexMethod.KNOWLEDGE_GRAPH],
@@ -56,7 +56,7 @@ def test_import_documents_from_files():
 
 def test_import_documents_from_datasource():
     kb = af.create_knowledge_base(
-        kb_id=uuid.UUID("01973588-65aa-4954-99fd-71eb5ecce167"),
+        id=uuid.UUID("01973588-65aa-4954-99fd-71eb5ecce167"),
         name="Test",
         description="This is a knowledge base for testing",
         index_methods=[IndexMethod.VECTOR_SEARCH],
@@ -72,7 +72,7 @@ def test_import_documents_from_datasource():
 
 def test_search_documents():
     kb = af.create_knowledge_base(
-        kb_id=uuid.UUID("01973588-65aa-4954-99fd-71eb5ecce167"),
+        id=uuid.UUID("01973588-65aa-4954-99fd-71eb5ecce167"),
         name="Test",
         description="This is a knowledge base for testing",
         index_methods=[IndexMethod.VECTOR_SEARCH],
@@ -89,7 +89,7 @@ def test_search_documents():
 
 def test_search_knowledge_graph():
     kb = af.create_knowledge_base(
-        kb_id=uuid.UUID("01973588-65aa-4954-99fd-71eb5ecce167"),
+        id=uuid.UUID("01973588-65aa-4954-99fd-71eb5ecce167"),
         name="Test",
         description="This is a knowledge base for testing",
         index_methods=[IndexMethod.KNOWLEDGE_GRAPH],