Merge pull request #28 from ChatBees/jluo-periodic-ingestion

new features and changes for release v1.1

Author: JuniusLuo

chatbees/client_models/chat.py

@@ -2,7 +2,7 @@
 
 from pydantic import BaseModel
 
-from chatbees.server_models.doc_api import AnswerReference
+from chatbees.server_models.doc_api import AskResponse
 from chatbees.utils.ask import ask
 from chatbees.utils.config import Config
 
@@ -17,17 +17,21 @@ class Chat(BaseModel):
     collection_name: str
     doc_name: Optional[str] = None
     history_messages: Optional[List[Tuple[str, str]]] = None
+    conversation_id: Optional[str] = None
 
-    def ask(self, question: str, top_k: int = 5) -> (str, List[AnswerReference]):
-        (answer, references) = ask(
+    def ask(self, question: str, top_k: int = 5) -> AskResponse:
+        resp = ask(
             Config.namespace,
             self.collection_name,
             question,
             top_k,
             doc_name=self.doc_name,
             history_messages=self.history_messages,
+            conversation_id=self.conversation_id,
         )
         if self.history_messages is None:
             self.history_messages = []
-        self.history_messages.append((question, answer))
-        return answer, references
+        self.history_messages.append((question, resp.answer))
+        if self.conversation_id is None:
+            self.conversation_id = resp.conversation_id
+        return resp

chatbees/client_models/collection.py

@@ -1,19 +1,20 @@
 import os
-from typing import List, Dict, Tuple, Any, Union
+from typing import List, Dict, Tuple, Any, Union, Optional
 from urllib import request
 
 from pydantic import BaseModel
 
 from chatbees.client_models.chat import Chat
 from chatbees.server_models.doc_api import (
     CrawlStatus,
-    IngestionStatus,
-    AnswerReference,
+    AskResponse,
     SearchReference,
 )
-from chatbees.server_models.chat import ConfigureChatRequest, ChatAttributes
+from chatbees.server_models.chat import ConfigureChatRequest
 from chatbees.server_models.ingestion_type import (
     IngestionType,
+    IngestionStatus,
+    ScheduleSpec,
     ConfluenceSpec,
     GDriveSpec,
     NotionSpec,
@@ -34,6 +35,8 @@
     DeleteCrawlRequest,
 )
 from chatbees.server_models.collection_api import (
+    ChatAttributes,
+    PeriodicIngest,
     DescribeCollectionResponse,
 )
 from chatbees.server_models.ingestion_api import (
@@ -43,8 +46,14 @@
     GetIngestionResponse,
     IndexIngestionRequest,
     DeleteIngestionRequest,
+    UpdatePeriodicIngestionRequest,
+    DeletePeriodicIngestionRequest,
 )
 from chatbees.server_models.search_api import SearchRequest, SearchResponse
+from chatbees.server_models.feedback_api import (
+    UnregisteredUser,
+    CreateOrUpdateFeedbackRequest,
+)
 from chatbees.utils.ask import ask
 from chatbees.utils.config import Config
 from chatbees.utils.file_upload import (
@@ -70,6 +79,10 @@ class Collection(BaseModel):
     # If true, collection can be read without an API key
     public_read: bool = False
 
+    chat_attributes: Optional[ChatAttributes] = None
+
+    periodic_ingests: Optional[List[PeriodicIngest]] = None
+
     def upload_document(self, path_or_url: str):
         """
         Uploads a local or web document into this collection.
@@ -146,7 +159,7 @@ def summarize_document(self, doc_name: str) -> str:
 
     def ask(
         self, question: str, top_k: int = 5, doc_name: str = None,
-    ) -> (str, List[AnswerReference]):
+    ) -> AskResponse:
         """
         Ask a question within the context of this collection.
 
@@ -205,7 +218,11 @@ def chat(self, doc_name: str = None) -> Chat:
             doc_name=doc_name
         )
 
-    def create_crawl(self, root_url: str, max_urls_to_crawl: int) -> str:
+    # TODO support update ScheduleSpec for periodic crawl.
+    # TODO deprecate crawl apis and switch to ingest apis.
+    def create_crawl(
+        self, root_url: str, max_urls_to_crawl: int, schedule: ScheduleSpec = None,
+    ) -> str:
         """
         Create a crawl task to crawl the root_url.
 
@@ -219,6 +236,7 @@ def create_crawl(self, root_url: str, max_urls_to_crawl: int) -> str:
             collection_name=self.name,
             root_url=root_url,
             max_urls_to_crawl=max_urls_to_crawl,
+            schedule=schedule,
         )
         resp = Config.post(url=url, data=req.model_dump_json())
         crawl_resp = CreateCrawlResponse.model_validate(resp.json())
@@ -227,7 +245,7 @@ def create_crawl(self, root_url: str, max_urls_to_crawl: int) -> str:
     def create_ingestion(
         self,
         ingestion_type: IngestionType,
-        ingestion_spec: Union[ConfluenceSpec, GDriveSpec, NotionSpec]
+        ingestion_spec: Union[ConfluenceSpec, GDriveSpec, NotionSpec],
     ) -> str:
         """
         Create an Ingestion task
@@ -249,6 +267,29 @@ def create_ingestion(
         ingest_resp = CreateIngestionResponse.model_validate(resp.json())
         return ingest_resp.ingestion_id
 
+    def update_periodic_ingestion(
+        self,
+        ingestion_type: IngestionType,
+        ingestion_spec: Union[ConfluenceSpec, GDriveSpec, NotionSpec]
+    ):
+        """
+        Update the periodic ingestion.
+
+        :param ingestion_type: the ingestion type
+        :param ingestion_spec: the spec for the ingestion. Currently, supports
+            - ConfluenceSpec
+            - GDriveSpec
+            - NotionSpec
+        :return: the id of the ingestion
+        """
+        url = f'{Config.get_base_url()}/docs/update_periodic_ingestion'
+        req = UpdatePeriodicIngestionRequest(
+            namespace_name=Config.namespace,
+            collection_name=self.name,
+            type=ingestion_type,
+            spec=ingestion_spec.model_dump())
+        Config.post(url=url, data=req.model_dump_json())
+
     def get_ingestion(self, ingestion_id: str) -> IngestionStatus:
         """
         Gets the Ingestion task status
@@ -294,6 +335,20 @@ def delete_ingestion(self, ingestion_type: IngestionType):
             type=ingestion_type)
         Config.post(url=url, data=req.model_dump_json())
 
+    def delete_periodic_ingestion(self, ingestion_type: IngestionType):
+        """
+        Delete the periodic ingestion for an ingestion type, e.g. a data source.
+        This does not delete the ingested data.
+
+        :param ingestion_type: the ingestion type
+        """
+        url = f'{Config.get_base_url()}/docs/delete_periodic_ingestion'
+        req = DeletePeriodicIngestionRequest(
+            namespace_name=Config.namespace,
+            collection_name=self.name,
+            type=ingestion_type)
+        Config.post(url=url, data=req.model_dump_json())
+
     def get_crawl(
         self, crawl_id: str,
     ) -> Tuple[CrawlStatus, Dict[str, PageStats]]:
@@ -373,6 +428,36 @@ def configure_chat(
         url = f'{Config.get_base_url()}/docs/configure_chat'
         Config.post(url=url, data=req.model_dump_json())
 
+        # update the local chat attributes
+        self.chat_attributes = req.chat_attributes
+
+    def create_or_update_feedback(
+        self,
+        request_id: str,
+        thumb_down: bool,
+        text_feedback: str = "",
+        unregistered_user: UnregisteredUser = None,
+    ):
+        """
+        Provides feedback for the ask or search.
+
+        :param request_id: the request_id of the ask or search
+        :param thumb_down: thumb up or down
+        :param text_feedback: optional text feedback
+        :param unregistered_user: optional information of the unregistered user
+        """
+        url = f'{Config.get_base_url()}/feedback/create_or_update'
+        req = CreateOrUpdateFeedbackRequest(
+            namespace_name=Config.namespace,
+            collection_name=self.name,
+            request_id=request_id,
+            thumb_down=thumb_down,
+            text_feedback=text_feedback,
+            unregistered_user=unregistered_user,
+        )
+        print(req.model_dump_json())
+        Config.post(url=url, data=req.model_dump_json())
+
 def describe_response_to_collection(
     collection_name: str,
     resp: DescribeCollectionResponse
@@ -386,5 +471,7 @@ def describe_response_to_collection(
     return Collection(
         name=collection_name,
         description=description,
-        public_read=public_read
+        public_read=public_read,
+        chat_attributes=resp.chat_attributes,
+        periodic_ingests=resp.periodic_ingests,
     )

chatbees/server_models/chat.py

@@ -2,24 +2,11 @@
 
 from pydantic import BaseModel
 
-from chatbees.server_models.collection_api import CollectionBaseRequest
-
-
-class ChatAttributes(BaseModel):
-    # Configure chatbot personality and style. For example:
-    #
-    # - a 1600s pirate, your name is 'Capitan Morgan'.
-    # - a helpful AI assistant
-    persona: Optional[str] = None
-
-    # Configure chatbot response when no relevant result is found. For example:
-    #
-    # - I do not have that information.
-    negative_response: Optional[str] = None
+from chatbees.server_models.collection_api import CollectionBaseRequest, ChatAttributes
 
 
 class ConfigureChatRequest(CollectionBaseRequest):
     """
     Configures a collection with custom q/a attributes
     """
-    chat_attributes: ChatAttributes
+    chat_attributes: ChatAttributes

chatbees/server_models/collection_api.py

@@ -1,7 +1,8 @@
-from typing import List, Optional
-
+from typing import Any, List, Optional
 from pydantic import BaseModel
 
+from chatbees.server_models.ingestion_type import IngestionType, IngestionStatus
+
 
 class CollectionBaseRequest(BaseModel):
     namespace_name: str
@@ -33,12 +34,40 @@ class ListCollectionsResponse(BaseModel):
     names: List[str]
 
 
+class ChatAttributes(BaseModel):
+    # Configure chatbot personality and style. For example:
+    #
+    # - a 1600s pirate, your name is 'Capitan Morgan'.
+    # - a helpful AI assistant
+    persona: Optional[str] = None
+
+    # Configure chatbot response when no relevant result is found. For example:
+    #
+    # - I do not have that information.
+    negative_response: Optional[str] = None
+
+
+class PeriodicIngest(BaseModel):
+    type: IngestionType
+    # the crawl spec
+    spec: Any
+    # the last ingestion epoch time
+    last_ingest_time: int
+    last_ingest_status: IngestionStatus
+
+
 class DescribeCollectionRequest(CollectionBaseRequest):
     pass
 
 
 class DescribeCollectionResponse(BaseModel):
     description: Optional[str] = None
 
+    # Chat attributes configured for this collection, if any
+    chat_attributes: Optional[ChatAttributes] = None
+
     # If true, the collection can be read without an API key
     public_read: Optional[bool] = None
+
+    # ingestins that run periodically for the collection
+    periodic_ingests: Optional[List[PeriodicIngest]] = None