.

Author: ssmith-pc

pinecone/repository/data/__init__.py

@@ -1,10 +1,10 @@
-from .repository import Repository, UpsertDocumentResponse
+from .repository import Repository
 
 
 _Repository = Repository  # alias for backwards compatibility
 
 
-__all__ = ["_Repository", "UpsertDocumentResponse"]
+__all__ = ["_Repository"]
 
 
 def __getattr__(name):

pinecone/repository/data/repository.bak

@@ -0,0 +1,89 @@
+import logging
+from typing import Optional, Dict, Any
+
+from pinecone.config import ConfigBuilder
+
+
+from pinecone.core.openapi.ckb_knowledge_data.api.document_operations_api import (
+    DocumentOperationsApi,
+)
+from pinecone.core.openapi.ckb_knowledge_data import API_VERSION
+
+from pinecone.openapi_support import ApiClient
+from pinecone.core.openapi.ckb_knowledge_data.models import (
+    UpsertDocumentResponse,
+    GetDocumentResponse,
+    ListDocumentsResponse,
+)
+
+from .interfaces import RepositoryInterface
+
+from pinecone.utils import setup_openapi_client
+
+from multiprocessing import cpu_count
+
+
+logger = logging.getLogger(__name__)
+""" :meta private: """
+
+
+class Repository(RepositoryInterface):
+    """
+    A client for interacting with a Pinecone Repository API.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        host: str,
+        pool_threads: Optional[int] = None,
+        additional_headers: Optional[Dict[str, str]] = {},
+        openapi_config=None,
+        **kwargs,
+    ):
+        self._config = ConfigBuilder.build(
+            api_key=api_key, host=host, additional_headers=additional_headers, **kwargs
+        )
+        """ :meta private: """
+        self._openapi_config = ConfigBuilder.build_openapi_config(self._config, openapi_config)
+        """ :meta private: """
+
+        if pool_threads is None:
+            self._pool_threads = 5 * cpu_count()
+            """ :meta private: """
+        else:
+            self._pool_threads = pool_threads
+            """ :meta private: """
+
+        if kwargs.get("connection_pool_maxsize", None):
+            self._openapi_config.connection_pool_maxsize = kwargs.get("connection_pool_maxsize")
+
+        self._repository_api = setup_openapi_client(
+            api_client_klass=ApiClient,
+            api_klass=DocumentOperationsApi,
+            config=self._config,
+            openapi_config=self._openapi_config,
+            pool_threads=self._pool_threads,
+            api_version=API_VERSION,
+        )
+
+        self._api_client = self._repository_api.api_client
+
+    def upsert(self, namespace: str, document: Dict[str, Any], **kwargs) -> UpsertDocumentResponse:
+        return self._repository_api.api_client.call_api(
+            resource_path=f"/ckb-stub-namespaces/{namespace}/documents/upsert",
+            method="POST",
+            path_params={"namespace": namespace},
+            body=document,
+            header_params={"Content-Type": "application/json", "Accept": "application/json"},
+            response_type=(UpsertDocumentResponse,),
+            _return_http_data_only=True,
+        )
+
+    def fetch(self, namespace: str, document_id: str, **kwargs) -> GetDocumentResponse:
+        return self._repository_api.get_document(
+            namespace=namespace, document_id=document_id, **kwargs
+        )
+
+    def list(self, namespace: str, **kwargs) -> ListDocumentsResponse:
+        return self._repository_api.list_documents(namespace=namespace, **kwargs)