letta-nightly 0.1.7.dev20240924104148__py3-none-any.whl

letta/cli/cli_load.py

@@ -0,0 +1,166 @@
+"""
+This file contains functions for loading data into Letta's archival storage.
+
+Data can be loaded with the following command, once a load function is defined:
+```
+letta load <data-connector-type> --name <dataset-name> [ADDITIONAL ARGS]
+```
+
+"""
+
+import uuid
+from typing import Annotated, List, Optional
+
+import typer
+
+from letta import create_client
+from letta.data_sources.connectors import DirectoryConnector
+
+app = typer.Typer()
+
+# NOTE: not supported due to llama-index breaking things (please reach out if you still need it)
+# @app.command("index")
+# def load_index(
+#    name: Annotated[str, typer.Option(help="Name of dataset to load.")],
+#    dir: Annotated[Optional[str], typer.Option(help="Path to directory containing index.")] = None,
+#    user_id: Annotated[Optional[uuid.UUID], typer.Option(help="User ID to associate with dataset.")] = None,
+# ):
+#    """Load a LlamaIndex saved VectorIndex into Letta"""
+#    if user_id is None:
+#        config = LettaConfig.load()
+#        user_id = uuid.UUID(config.anon_clientid)
+#
+#    try:
+#        # load index data
+#        storage_context = StorageContext.from_defaults(persist_dir=dir)
+#        loaded_index = load_index_from_storage(storage_context)
+#
+#        # hacky code to extract out passages/embeddings (thanks a lot, llama index)
+#        embed_dict = loaded_index._vector_store._data.embedding_dict
+#        node_dict = loaded_index._docstore.docs
+#
+#        # create storage connector
+#        config = LettaConfig.load()
+#        if user_id is None:
+#            user_id = uuid.UUID(config.anon_clientid)
+#
+#        passages = []
+#        for node_id, node in node_dict.items():
+#            vector = embed_dict[node_id]
+#            node.embedding = vector
+#            # assume embedding are the same as config
+#            passages.append(
+#                Passage(
+#                    text=node.text,
+#                    embedding=np.array(vector),
+#                    embedding_dim=config.default_embedding_config.embedding_dim,
+#                    embedding_model=config.default_embedding_config.embedding_model,
+#                )
+#            )
+#            assert config.default_embedding_config.embedding_dim == len(
+#                vector
+#            ), f"Expected embedding dimension {config.default_embedding_config.embedding_dim}, got {len(vector)}"
+#
+#        if len(passages) == 0:
+#            raise ValueError(f"No passages found in index {dir}")
+#
+#        insert_passages_into_source(passages, name, user_id, config)
+#    except ValueError as e:
+#        typer.secho(f"Failed to load index from provided information.\n{e}", fg=typer.colors.RED)
+
+
+default_extensions = ".txt,.md,.pdf"
+
+
+@app.command("directory")
+def load_directory(
+    name: Annotated[str, typer.Option(help="Name of dataset to load.")],
+    input_dir: Annotated[Optional[str], typer.Option(help="Path to directory containing dataset.")] = None,
+    input_files: Annotated[List[str], typer.Option(help="List of paths to files containing dataset.")] = [],
+    recursive: Annotated[bool, typer.Option(help="Recursively search for files in directory.")] = False,
+    extensions: Annotated[str, typer.Option(help="Comma separated list of file extensions to load")] = default_extensions,
+    user_id: Annotated[Optional[uuid.UUID], typer.Option(help="User ID to associate with dataset.")] = None,  # TODO: remove
+    description: Annotated[Optional[str], typer.Option(help="Description of the source.")] = None,
+):
+    client = create_client()
+
+    # create connector
+    connector = DirectoryConnector(input_files=input_files, input_directory=input_dir, recursive=recursive, extensions=extensions)
+
+    # create source
+    source = client.create_source(name=name)
+
+    # load data
+    try:
+        client.load_data(connector, source_name=name)
+    except Exception as e:
+        typer.secho(f"Failed to load data from provided information.\n{e}", fg=typer.colors.RED)
+        client.delete_source(source.id)
+
+
+# @app.command("webpage")
+# def load_webpage(
+#    name: Annotated[str, typer.Option(help="Name of dataset to load.")],
+#    urls: Annotated[List[str], typer.Option(help="List of urls to load.")],
+# ):
+#    try:
+#        from llama_index.readers.web import SimpleWebPageReader
+#
+#        docs = SimpleWebPageReader(html_to_text=True).load_data(urls)
+#        store_docs(name, docs)
+#
+#    except ValueError as e:
+#        typer.secho(f"Failed to load webpage from provided information.\n{e}", fg=typer.colors.RED)
+
+
+@app.command("vector-database")
+def load_vector_database(
+    name: Annotated[str, typer.Option(help="Name of dataset to load.")],
+    uri: Annotated[str, typer.Option(help="Database URI.")],
+    table_name: Annotated[str, typer.Option(help="Name of table containing data.")],
+    text_column: Annotated[str, typer.Option(help="Name of column containing text.")],
+    embedding_column: Annotated[str, typer.Option(help="Name of column containing embedding.")],
+    user_id: Annotated[Optional[uuid.UUID], typer.Option(help="User ID to associate with dataset.")] = None,
+):
+    """Load pre-computed embeddings into Letta from a database."""
+    raise NotImplementedError
+    # try:
+    #    config = LettaConfig.load()
+    #    connector = VectorDBConnector(
+    #        uri=uri,
+    #        table_name=table_name,
+    #        text_column=text_column,
+    #        embedding_column=embedding_column,
+    #        embedding_dim=config.default_embedding_config.embedding_dim,
+    #    )
+    #    if not user_id:
+    #        user_id = uuid.UUID(config.anon_clientid)
+
+    #    ms = MetadataStore(config)
+    #    source = Source(
+    #        name=name,
+    #        user_id=user_id,
+    #        embedding_model=config.default_embedding_config.embedding_model,
+    #        embedding_dim=config.default_embedding_config.embedding_dim,
+    #    )
+    #    ms.create_source(source)
+    #    passage_storage = StorageConnector.get_storage_connector(TableType.PASSAGES, config, user_id)
+    #    # TODO: also get document store
+
+    #    # ingest data into passage/document store
+    #    try:
+    #        num_passages, num_documents = load_data(
+    #            connector=connector,
+    #            source=source,
+    #            embedding_config=config.default_embedding_config,
+    #            document_store=None,
+    #            passage_store=passage_storage,
+    #        )
+    #        print(f"Loaded {num_passages} passages and {num_documents} documents from {name}")
+    #    except Exception as e:
+    #        typer.secho(f"Failed to load data from provided information.\n{e}", fg=typer.colors.RED)
+    #        ms.delete_source(source_id=source.id)
+
+    # except ValueError as e:
+    #    typer.secho(f"Failed to load VectorDB from provided information.\n{e}", fg=typer.colors.RED)
+    #    raise

letta/client/admin.py

@@ -0,0 +1,171 @@
+from typing import List, Optional
+
+import requests
+from requests import HTTPError
+
+from letta.functions.functions import parse_source_code
+from letta.functions.schema_generator import generate_schema
+from letta.schemas.api_key import APIKey, APIKeyCreate
+from letta.schemas.organization import Organization, OrganizationCreate
+from letta.schemas.user import User, UserCreate
+
+
+class Admin:
+    """
+    Admin client allows admin-level operations on the Letta server.
+    - Creating users
+    - Generating user keys
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        token: str,
+        api_prefix: str = "v1",
+    ):
+        self.base_url = base_url
+        self.api_prefix = api_prefix
+        self.token = token
+        self.headers = {"accept": "application/json", "content-type": "application/json", "authorization": f"Bearer {token}"}
+
+    def get_users(self, cursor: Optional[str] = None, limit: Optional[int] = 50) -> List[User]:
+        params = {}
+        if cursor:
+            params["cursor"] = str(cursor)
+        if limit:
+            params["limit"] = limit
+        response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/users", params=params, headers=self.headers)
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        return [User(**user) for user in response.json()]
+
+    def create_key(self, user_id: str, key_name: Optional[str] = None) -> APIKey:
+        request = APIKeyCreate(user_id=user_id, name=key_name)
+        response = requests.post(f"{self.base_url}/{self.api_prefix}/admin/users/keys", headers=self.headers, json=request.model_dump())
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        return APIKey(**response.json())
+
+    def get_keys(self, user_id: str) -> List[APIKey]:
+        params = {"user_id": str(user_id)}
+        response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/users/keys", params=params, headers=self.headers)
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        return [APIKey(**key) for key in response.json()]
+
+    def delete_key(self, api_key: str) -> APIKey:
+        params = {"api_key": api_key}
+        response = requests.delete(f"{self.base_url}/{self.api_prefix}/admin/users/keys", params=params, headers=self.headers)
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        return APIKey(**response.json())
+
+    def create_user(self, name: Optional[str] = None, org_id: Optional[str] = None) -> User:
+        request = UserCreate(name=name, org_id=org_id)
+        response = requests.post(f"{self.base_url}/{self.api_prefix}/admin/users", headers=self.headers, json=request.model_dump())
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        response_json = response.json()
+        return User(**response_json)
+
+    def delete_user(self, user_id: str) -> User:
+        params = {"user_id": str(user_id)}
+        response = requests.delete(f"{self.base_url}/{self.api_prefix}/admin/users", params=params, headers=self.headers)
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        return User(**response.json())
+
+    def create_organization(self, name: Optional[str] = None) -> Organization:
+        request = OrganizationCreate(name=name)
+        response = requests.post(f"{self.base_url}/{self.api_prefix}/admin/orgs", headers=self.headers, json=request.model_dump())
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        response_json = response.json()
+        return Organization(**response_json)
+
+    def delete_organization(self, org_id: str) -> Organization:
+        params = {"org_id": str(org_id)}
+        response = requests.delete(f"{self.base_url}/{self.api_prefix}/admin/orgs", params=params, headers=self.headers)
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        return Organization(**response.json())
+
+    def get_organizations(self, cursor: Optional[str] = None, limit: Optional[int] = 50) -> List[Organization]:
+        params = {}
+        if cursor:
+            params["cursor"] = str(cursor)
+        if limit:
+            params["limit"] = limit
+        response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/orgs", params=params, headers=self.headers)
+        if response.status_code != 200:
+            raise HTTPError(response.json())
+        return [Organization(**org) for org in response.json()]
+
+    def _reset_server(self):
+        # DANGER: this will delete all users and keys
+        # clear all state associated with users
+        # TODO: clear out all agents, presets, etc.
+        users = self.get_users()
+        for user in users:
+            keys = self.get_keys(user.id)
+            for key in keys:
+                self.delete_key(key.key)
+            self.delete_user(user.id)
+
+    # tools
+    def create_tool(
+        self,
+        func,
+        name: Optional[str] = None,
+        update: Optional[bool] = True,  # TODO: actually use this
+        tags: Optional[List[str]] = None,
+    ):
+        """Create a tool
+        Args:
+            func (callable): The function to create a tool for.
+            tags (Optional[List[str]], optional): Tags for the tool. Defaults to None.
+            update (bool, optional): Update the tool if it already exists. Defaults to True.
+        Returns:
+            Tool object
+        """
+
+        # TODO: check if tool already exists
+        # TODO: how to load modules?
+        # parse source code/schema
+        source_code = parse_source_code(func)
+        json_schema = generate_schema(func, name)
+        source_type = "python"
+        json_schema["name"]
+
+        if "memory" in tags:
+            # special modifications to memory functions
+            # self.memory -> self.memory.memory, since Agent.memory.memory needs to be modified (not BaseMemory.memory)
+            source_code = source_code.replace("self.memory", "self.memory.memory")
+
+        # create data
+        data = {"source_code": source_code, "source_type": source_type, "tags": tags, "json_schema": json_schema}
+        CreateToolRequest(**data)  # validate
+
+        # make REST request
+        response = requests.post(f"{self.base_url}/{self.api_prefix}/admin/tools", json=data, headers=self.headers)
+        if response.status_code != 200:
+            raise ValueError(f"Failed to create tool: {response.text}")
+        return ToolModel(**response.json())
+
+    def list_tools(self):
+        response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/tools", headers=self.headers)
+        return ListToolsResponse(**response.json()).tools
+
+    def delete_tool(self, name: str):
+        response = requests.delete(f"{self.base_url}/{self.api_prefix}/admin/tools/{name}", headers=self.headers)
+        if response.status_code != 200:
+            raise ValueError(f"Failed to delete tool: {response.text}")
+        return response.json()
+
+    def get_tool(self, name: str):
+        response = requests.get(f"{self.base_url}/{self.api_prefix}/admin/tools/{name}", headers=self.headers)
+        if response.status_code == 404:
+            return None
+        elif response.status_code != 200:
+            raise ValueError(f"Failed to get tool: {response.text}")
+        return ToolModel(**response.json())