From ecda804ca82b670da6c1cd6c7506070b6d70f960 Mon Sep 17 00:00:00 2001
From: tsukasa <puotosssa@live.com>
Date: Thu, 28 Sep 2023 11:05:24 +0800
Subject: [PATCH] remove /knowledge query command

---
 src/aios_kernel/knowledge_base.py    | 17 ++++++++++-------
 src/service/aios_shell/aios_shell.py | 13 ++-----------
 2 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py
index 2112b90..8e4de5a 100644
--- a/src/aios_kernel/knowledge_base.py
+++ b/src/aios_kernel/knowledge_base.py
@@ -163,9 +163,9 @@ class KnowledgeBase:
         self.store.get_object_store().put_object(object.calculate_id(), object.encode())
         await self.__do_embedding(object)
     
-    async def query_objects(self, tokens: str) -> [ObjectID]:
+    async def query_objects(self, tokens: str, topk: int) -> [ObjectID]:
         vector = await self.compute_kernel.do_text_embedding(tokens, self._default_text_model)
-        return await self.store.get_vector_store(self._default_text_model).query(vector, 10)
+        return await self.store.get_vector_store(self._default_text_model).query(vector, topk)
 
     def __load_object(self, object_id: ObjectID) -> KnowledgeObject:
         if object_id.get_object_type() == ObjectType.Document:
@@ -193,8 +193,7 @@ class KnowledgeBase:
                 results[str(root_object_id)].append(object_id)
             else:
                 results[str(root_object_id)] = [root_object_id, object_id]
-
-        content = "*** I have provided the following known information for your reference with json format:\n"
+        content = ""
         result_desc = []
         for result in results.values():
             # first element in result is the root object
@@ -244,6 +243,10 @@ class KnowledgeEnvironment(Environment):
                                             self._query, 
                                             query_param))
         
-    async def _query(tokens: str, index: int):
-        object_ids = await KnowledgeBase().query_objects(tokens)
-        KnowledgeBase().tokens_from_objects(object_ids)
\ No newline at end of file
+    async def _query(self, tokens: str, index: int=0):
+        object_ids = await KnowledgeBase().query_objects(tokens, 4)
+        if len(object_ids) <= index:
+            return "*** I have no more information for your reference.\n"
+        else:
+            content = "*** I have provided the following known information for your reference with json format:\n"
+            return content + KnowledgeBase().tokens_from_objects(object_ids[index:index + 1])
\ No newline at end of file
diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py
index d15c15f..7b5f654 100644
--- a/src/service/aios_shell/aios_shell.py
+++ b/src/service/aios_shell/aios_shell.py
@@ -320,8 +320,7 @@ class AIOS_Shell:
     async def handle_knowledge_commands(self, args):
         show_text = FormattedText([("class:title", "sub command not support!\n" 
                               "/knowledge add email | dir\n"
-                              "/knowledge journal [$topn]\n"
-                              "/knowledge query $query\n")])
+                              "/knowledge journal [$topn]\n")])
         if len(args) < 1:
             return show_text
         sub_cmd = args[0]
@@ -358,13 +357,6 @@ class AIOS_Shell:
             topn = 10 if len(args) == 1 else int(args[1])
             journals = [str(journal) for journal in KnowledgePipline.get_instance().get_latest_journals(topn)]
             print_formatted_text("\r\n".join(journals))
-        if sub_cmd == "query":
-            if len(args) < 2:
-                return show_text
-            prompt = AgentPrompt()
-            prompt.messages.append({"role": "user", "content":" ".join(args[1:])})
-            result = await KnowledgeBase().query_prompt(prompt)
-            print_formatted_text(result.as_str())
 
     async def call_func(self,func_name, args):
         match func_name:
@@ -674,8 +666,7 @@ async def main():
                                '/connect $target',
                                '/contact $name',
                                '/knowledge add email | dir',
-                               '/knowledge journal [$topn]',
-                               '/knowledge query $query' 
+                               '/knowledge journal [$topn]', 
                                '/set_config $key',
                                '/enable $feature',
                                '/disable $feature',