Local fs testing

chatbees/__init__.py

@@ -1,5 +1,6 @@
 from .client.collection_management import *
 from .client.admin_management import *
+from .client.application_management import *
 from .client_models.collection import *
 from .client_models.chat import *
 from .server_models.doc_api import *

chatbees/client/admin_management.py

@@ -1,13 +1,17 @@
 from typing import List
-from chatbees.server_models.admin_api import CreateApiKeyRequest, CreateApiKeyResponse
+from chatbees.server_models.admin_api import (
+    CreateApiKeyRequest,
+    CreateApiKeyResponse,
+    EmailAccountRequest, AccountLoginResponse,
+)
 from chatbees.server_models.ingestion_api import (
     ConnectorReference,
     ListConnectorsRequest,
     ListConnectorsResponse,
 )
 from chatbees.utils.config import Config
 
-__all__ = ["init", "list_connectors"]
+__all__ = ["init", "list_connectors", "email_login"]
 
 
 def init(
@@ -30,6 +34,33 @@ def init(
     Config.namespace = namespace
     Config.validate_setup()
 
+def email_login(
+    email: str,
+    password: str,
+):
+    """
+    Initialize the ChatBees client.
+
+    Args:
+        api_key (str): The API key to authenticate requests.
+        account_id (str): The account ID.
+        namespace (str, optional): The namespace to use.
+    Raises:
+        ValueError: If the provided config is invalid
+    """
+    url = f"{Config.get_base_url()}/account/email"
+    req = EmailAccountRequest(email=email, password=password)
+    resp = Config.post(
+        url=url,
+        data=req.model_dump_json(),
+        enforce_api_key=False,
+    )
+    resp = AccountLoginResponse.model_validate(resp.json())
+
+    Config.api_key = resp.shortlive_api_key
+    Config.account_id = resp.account_id
+    Config.namespace = resp.default_namespace
+    Config.validate_setup()
 
 def list_connectors() -> List[ConnectorReference]:
     url = f'{Config.get_base_url()}/connectors/list'

chatbees/client/application_management.py

@@ -0,0 +1,85 @@
+__all__ = ["create_gpt_application", "create_collection_application", "delete_application", "list_applications"]
+
+from typing import Optional, List
+
+from chatbees.server_models.application import (
+    Application,
+    ApplicationType,
+    CollectionTarget,
+    GPTTarget,
+)
+from chatbees.server_models.application_api import (
+    CreateApplicationRequest,
+    DeleteApplicationRequest, ListApplicationsResponse, ListApplicationsRequest,
+)
+from chatbees.server_models.collection_api import ChatAttributes
+from chatbees.utils.config import Config
+
+def create_gpt_application(
+    application_name: str,
+    provider: str,
+    model: str,
+    description: Optional[str] = None,
+    chat_attrs: Optional[ChatAttributes] = None
+) -> Application:
+    """
+    Create a new collection application in ChatBees.
+
+    """
+    url = f'{Config.get_base_url()}/applications/create'
+    application = Application(
+        application_name=application_name,
+        application_desc=description,
+        application_type=ApplicationType.GPT,
+        chat_attrs=chat_attrs,
+        application_target=GPTTarget(
+            provider=provider, model=model).model_dump_json())
+    req = CreateApplicationRequest(namespace_name=Config.namespace, application=application)
+    Config.post(url=url, data=req.model_dump_json())
+    return application
+
+def create_collection_application(
+    application_name: str,
+    collection_name: str,
+    description: Optional[str] = None,
+    chat_attrs: Optional[ChatAttributes] = None
+) -> Application:
+    """
+    Create a new collection application in ChatBees.
+
+    """
+    url = f'{Config.get_base_url()}/applications/create'
+    application = Application(
+        application_name=application_name,
+        application_desc=description,
+        application_type=ApplicationType.COLLECTION,
+        chat_attrs=chat_attrs,
+        application_target=CollectionTarget(collection_name=collection_name).model_dump_json())
+    req = CreateApplicationRequest(namespace_name=Config.namespace, application=application)
+    Config.post(url=url, data=req.model_dump_json())
+    return application
+
+
+def delete_application(application_name: str):
+    """
+    Deletes an application
+
+    Args:
+        application_name (str): The name of the application.
+    """
+    url = f'{Config.get_base_url()}/applications/delete'
+    req = DeleteApplicationRequest(namespace_name=Config.namespace, application_name=application_name)
+    Config.post(url=url, data=req.model_dump_json())
+
+
+def list_applications() -> List[Application]:
+    """
+    List all applications in account.
+
+    Returns:
+        List[Application]: A list of application objects.
+    """
+    url = f'{Config.get_base_url()}/applications/list'
+    req = ListApplicationsRequest(namespace_name=Config.namespace)
+    resp = Config.post(url=url, data=req.model_dump_json())
+    return ListApplicationsResponse.model_validate(resp.json()).applications

chatbees/client_models/chat.py

@@ -2,8 +2,13 @@
 
 from pydantic import BaseModel
 
+from chatbees.server_models.conversation import (
+    ConversationMeta,
+    ListConversationsRequest,
+    ListConversationsResponse, GetConversationRequest, GetConversationResponse,
+)
 from chatbees.server_models.doc_api import AskResponse
-from chatbees.utils.ask import ask
+from chatbees.utils.ask import ask, ask_application
 from chatbees.utils.config import Config
 
 __all__ = ["Chat"]
@@ -13,22 +18,82 @@ class Chat(BaseModel):
     """
     A new chatbot instance that supports conversational Q and A.
     """
-    namespace_name: str
-    collection_name: str
+    namespace_name: Optional[str] = None
+    collection_name: Optional[str] = None
+    application_name: Optional[str] = None
     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) -> 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,
+    def _validate_names(self):
+        if self.application_name is not None and self.namespace_name is None and self.collection_name is None:
+            return
+        if self.application_name is None and self.namespace_name is not None and self.collection_name is not None:
+            return
+        raise ValueError(f"Chat must specify either both namespace and collection, or application")
+
+    @classmethod
+    def list_conversations(cls, collection_name=None, application_name=None) -> List[ConversationMeta]:
+        url = f'{Config.get_base_url()}/conversations/list'
+        namespace = Config.namespace if collection_name is not None else None
+        req = ListConversationsRequest(
+            namespace_name=namespace,
+            collection_name=collection_name,
+            application_name=application_name
+        )
+        resp = Config.post(
+            url=url,
+            data=req.model_dump_json(),
+        )
+        return ListConversationsResponse.model_validate(resp.json()).conversations
+
+    @classmethod
+    def from_conversation(cls, conversation_id, collection_name=None, application_name=None) -> "Chat":
+        url = f'{Config.get_base_url()}/conversations/get'
+        req = GetConversationRequest(
+            conversation_id=conversation_id,
+        )
+        resp = Config.post(
+            url=url,
+            data=req.model_dump_json(),
         )
+        chat = Chat()
+        convos = GetConversationResponse.model_validate(resp.json()).conversation
+        chat.conversation_id = conversation_id
+        chat.collection_name = collection_name
+        chat.application_name = application_name
+
+        # TODO: Fix this convoluted logic...openai accepts a plain array of
+        #  {role, content}. Simply preserve this and pass it over to openai
+        paired_content = [(convos.messages[i].content, convos.messages[i + 1].content) for i in
+                          range(0, len(convos.messages) - 1, 2)]
+        chat.history_messages = paired_content
+        return chat
+
+
+    def ask(self, question: str, top_k: int = 5) -> AskResponse:
+        self._validate_names()
+
+        if self.application_name is None:
+            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,
+            )
+        else:
+            resp = ask_application(
+                self.application_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, resp.answer))

chatbees/client_models/collection.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from typing import List, Dict, Tuple, Any, Union, Optional
 from urllib import request
@@ -8,7 +9,7 @@
 from chatbees.server_models.doc_api import (
     CrawlStatus,
     AskResponse,
-    SearchReference,
+    SearchReference, GetDocRequest, AsyncAddDocResponse, PendingDocumentMetadata,
 )
 from chatbees.server_models.chat import ConfigureChatRequest
 from chatbees.server_models.ingestion_type import (
@@ -90,6 +91,30 @@ class Collection(BaseModel):
 
     periodic_ingests: Optional[List[PeriodicIngest]] = None
 
+
+    def download_document(self, doc_name: str, save_path = '.'):
+        """
+        Uploads a local or web document into this collection.
+
+        :param path_or_url: Local file path or the URL of a document. URL must
+                            contain scheme (http or https) prefix.
+        :return:
+        """
+        url = f'{Config.get_base_url()}/docs/get'
+        req = GetDocRequest(namespace_name=Config.namespace,
+                            collection_name=self.name, doc_name=doc_name)
+
+        response = Config.post(url=url, data=req.model_dump_json())
+        response.raise_for_status()
+        print(f'GET file response is {response}')
+
+        # Write the file to the local save path
+        filename = os.path.join(save_path, doc_name)
+        with open(filename, 'wb') as file:
+            file.write(response.content)
+        print(f"Document '{doc_name}' downloaded successfully to '{save_path}'.")
+        return filename
+
     def upload_document(self, path_or_url: str):
         """
         Uploads a local or web document into this collection.
@@ -118,6 +143,41 @@ def upload_document(self, path_or_url: str):
                     url=url, files={'file': (fname, f)},
                     data={'request': req.model_dump_json()})
 
+    def upload_documents_async(self, path_or_urls: List[str]) -> List[str]:
+        """
+        Uploads a local or web document into this collection.
+
+        :param path_or_urls: Local file path or the URL of a document. URL must
+                             contain scheme (http or https) prefix.
+        :return:
+        """
+        url = f'{Config.get_base_url()}/docs/add_async'
+        req = AddDocRequest(namespace_name=Config.namespace,
+                            collection_name=self.name)
+        #files = [('files', open(f"{current_dir}/{fname}", "rb")) for fname in fnames]
+        files = []
+
+        for path_or_url in path_or_urls:
+            if is_url(path_or_url):
+                validate_url_file(path_or_url)
+                files.append(('files', request.urlopen(path_or_url) ))
+            else:
+                # Handle tilde "~/blah"
+                path_or_url = os.path.expanduser(path_or_url)
+                validate_file(path_or_url)
+                files.append(('files', open(path_or_url, 'rb')))
+        resp = Config.post(
+            url=url, files=files,
+            data={'request': req.model_dump_json()})
+        for _, file in files:
+            try:
+                file.close()
+                logging.info(f"Closed file")
+            except Exception:
+                pass
+        add_resp = AsyncAddDocResponse.model_validate(resp.json())
+        return add_resp.task_ids
+
     def delete_document(self, doc_name: str):
         """
         Deletes the document.
@@ -147,6 +207,21 @@ def list_documents(self) -> List[str]:
         list_resp = ListDocsResponse.model_validate(resp.json())
         return list_resp.doc_names
 
+    def list_pending_documents(self) -> List[PendingDocumentMetadata]:
+        """
+        List the documents.
+
+        :return: A list of the documents
+        """
+        url = f'{Config.get_base_url()}/docs/list'
+        req = ListDocsRequest(
+            namespace_name=Config.namespace,
+            collection_name=self.name,
+        )
+        resp = Config.post(url=url, data=req.model_dump_json())
+        list_resp = ListDocsResponse.model_validate(resp.json())
+        return list_resp.pending_documents
+
     def summarize_document(self, doc_name: str) -> str:
         """
         Returns a summary of the document.