llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/adapters/memory/pgvector/pgvector.py

@@ -1,211 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from typing import List, Tuple
-
-import psycopg2
-from numpy.typing import NDArray
-from psycopg2 import sql
-from psycopg2.extras import execute_values, Json
-
-from pydantic import BaseModel, parse_obj_as
-
-from llama_stack.apis.memory import *  # noqa: F403
-
-from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
-from llama_stack.providers.utils.memory.vector_store import (
-    ALL_MINILM_L6_V2_DIMENSION,
-    BankWithIndex,
-    EmbeddingIndex,
-)
-
-from .config import PGVectorConfig
-
-
-def check_extension_version(cur):
-    cur.execute("SELECT extversion FROM pg_extension WHERE extname = 'vector'")
-    result = cur.fetchone()
-    return result[0] if result else None
-
-
-def upsert_models(cur, keys_models: List[Tuple[str, BaseModel]]):
-    query = sql.SQL(
-        """
-        INSERT INTO metadata_store (key, data)
-        VALUES %s
-        ON CONFLICT (key) DO UPDATE
-        SET data = EXCLUDED.data
-    """
-    )
-
-    values = [(key, Json(model.dict())) for key, model in keys_models]
-    execute_values(cur, query, values, template="(%s, %s)")
-
-
-def load_models(cur, cls):
-    query = "SELECT key, data FROM metadata_store"
-    cur.execute(query)
-    rows = cur.fetchall()
-    return [parse_obj_as(cls, row["data"]) for row in rows]
-
-
-class PGVectorIndex(EmbeddingIndex):
-    def __init__(self, bank: MemoryBankDef, dimension: int, cursor):
-        self.cursor = cursor
-        self.table_name = f"vector_store_{bank.identifier}"
-
-        self.cursor.execute(
-            f"""
-            CREATE TABLE IF NOT EXISTS {self.table_name} (
-                id TEXT PRIMARY KEY,
-                document JSONB,
-                embedding vector({dimension})
-            )
-        """
-        )
-
-    async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
-        assert len(chunks) == len(
-            embeddings
-        ), f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
-
-        values = []
-        for i, chunk in enumerate(chunks):
-            values.append(
-                (
-                    f"{chunk.document_id}:chunk-{i}",
-                    Json(chunk.dict()),
-                    embeddings[i].tolist(),
-                )
-            )
-
-        query = sql.SQL(
-            f"""
-        INSERT INTO {self.table_name} (id, document, embedding)
-        VALUES %s
-        ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding, document = EXCLUDED.document
-    """
-        )
-        execute_values(self.cursor, query, values, template="(%s, %s, %s::vector)")
-
-    async def query(self, embedding: NDArray, k: int) -> QueryDocumentsResponse:
-        self.cursor.execute(
-            f"""
-        SELECT document, embedding <-> %s::vector AS distance
-        FROM {self.table_name}
-        ORDER BY distance
-        LIMIT %s
-    """,
-            (embedding.tolist(), k),
-        )
-        results = self.cursor.fetchall()
-
-        chunks = []
-        scores = []
-        for doc, dist in results:
-            chunks.append(Chunk(**doc))
-            scores.append(1.0 / float(dist))
-
-        return QueryDocumentsResponse(chunks=chunks, scores=scores)
-
-
-class PGVectorMemoryAdapter(Memory, MemoryBanksProtocolPrivate):
-    def __init__(self, config: PGVectorConfig) -> None:
-        print(f"Initializing PGVectorMemoryAdapter -> {config.host}:{config.port}")
-        self.config = config
-        self.cursor = None
-        self.conn = None
-        self.cache = {}
-
-    async def initialize(self) -> None:
-        try:
-            self.conn = psycopg2.connect(
-                host=self.config.host,
-                port=self.config.port,
-                database=self.config.db,
-                user=self.config.user,
-                password=self.config.password,
-            )
-            self.cursor = self.conn.cursor()
-
-            version = check_extension_version(self.cursor)
-            if version:
-                print(f"Vector extension version: {version}")
-            else:
-                raise RuntimeError("Vector extension is not installed.")
-
-            self.cursor.execute(
-                """
-                CREATE TABLE IF NOT EXISTS metadata_store (
-                    key TEXT PRIMARY KEY,
-                    data JSONB
-                )
-            """
-            )
-        except Exception as e:
-            import traceback
-
-            traceback.print_exc()
-            raise RuntimeError("Could not connect to PGVector database server") from e
-
-    async def shutdown(self) -> None:
-        pass
-
-    async def register_memory_bank(
-        self,
-        memory_bank: MemoryBankDef,
-    ) -> None:
-        assert (
-            memory_bank.type == MemoryBankType.vector.value
-        ), f"Only vector banks are supported {memory_bank.type}"
-
-        upsert_models(
-            self.cursor,
-            [
-                (memory_bank.identifier, memory_bank),
-            ],
-        )
-
-        index = BankWithIndex(
-            bank=memory_bank,
-            index=PGVectorIndex(memory_bank, ALL_MINILM_L6_V2_DIMENSION, self.cursor),
-        )
-        self.cache[memory_bank.identifier] = index
-
-    async def list_memory_banks(self) -> List[MemoryBankDef]:
-        banks = load_models(self.cursor, MemoryBankDef)
-        for bank in banks:
-            if bank.identifier not in self.cache:
-                index = BankWithIndex(
-                    bank=bank,
-                    index=PGVectorIndex(bank, ALL_MINILM_L6_V2_DIMENSION, self.cursor),
-                )
-                self.cache[bank.identifier] = index
-        return banks
-
-    async def insert_documents(
-        self,
-        bank_id: str,
-        documents: List[MemoryBankDocument],
-        ttl_seconds: Optional[int] = None,
-    ) -> None:
-        index = self.cache.get(bank_id, None)
-        if not index:
-            raise ValueError(f"Bank {bank_id} not found")
-
-        await index.insert_documents(documents)
-
-    async def query_documents(
-        self,
-        bank_id: str,
-        query: InterleavedTextMedia,
-        params: Optional[Dict[str, Any]] = None,
-    ) -> QueryDocumentsResponse:
-        index = self.cache.get(bank_id, None)
-        if not index:
-            raise ValueError(f"Bank {bank_id} not found")
-
-        return await index.query_documents(query, params)

llama_stack/providers/adapters/memory/sample/sample.py

@@ -1,23 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from .config import SampleConfig
-
-
-from llama_stack.apis.memory import *  # noqa: F403
-
-
-class SampleMemoryImpl(Memory):
-    def __init__(self, config: SampleConfig):
-        self.config = config
-
-    async def register_memory_bank(self, memory_bank: MemoryBankDef) -> None:
-        # these are the memory banks the Llama Stack will use to route requests to this provider
-        # perform validation here if necessary
-        pass
-
-    async def initialize(self):
-        pass

llama_stack/providers/adapters/memory/weaviate/__init__.py

@@ -1,15 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from .config import WeaviateConfig, WeaviateRequestProviderData  # noqa: F401
-
-
-async def get_adapter_impl(config: WeaviateConfig, _deps):
-    from .weaviate import WeaviateMemoryAdapter
-
-    impl = WeaviateMemoryAdapter(config)
-    await impl.initialize()
-    return impl

llama_stack/providers/adapters/memory/weaviate/weaviate.py

@@ -1,190 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-import json
-
-from typing import Any, Dict, List, Optional
-
-import weaviate
-import weaviate.classes as wvc
-from numpy.typing import NDArray
-from weaviate.classes.init import Auth
-
-from llama_stack.apis.memory import *  # noqa: F403
-from llama_stack.distribution.request_headers import NeedsRequestProviderData
-from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
-from llama_stack.providers.utils.memory.vector_store import (
-    BankWithIndex,
-    EmbeddingIndex,
-)
-
-from .config import WeaviateConfig, WeaviateRequestProviderData
-
-
-class WeaviateIndex(EmbeddingIndex):
-    def __init__(self, client: weaviate.Client, collection_name: str):
-        self.client = client
-        self.collection_name = collection_name
-
-    async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
-        assert len(chunks) == len(
-            embeddings
-        ), f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
-
-        data_objects = []
-        for i, chunk in enumerate(chunks):
-            data_objects.append(
-                wvc.data.DataObject(
-                    properties={
-                        "chunk_content": chunk.json(),
-                    },
-                    vector=embeddings[i].tolist(),
-                )
-            )
-
-        # Inserting chunks into a prespecified Weaviate collection
-        collection = self.client.collections.get(self.collection_name)
-
-        # TODO: make this async friendly
-        collection.data.insert_many(data_objects)
-
-    async def query(self, embedding: NDArray, k: int) -> QueryDocumentsResponse:
-        collection = self.client.collections.get(self.collection_name)
-
-        results = collection.query.near_vector(
-            near_vector=embedding.tolist(),
-            limit=k,
-            return_metadata=wvc.query.MetadataQuery(distance=True),
-        )
-
-        chunks = []
-        scores = []
-        for doc in results.objects:
-            chunk_json = doc.properties["chunk_content"]
-            try:
-                chunk_dict = json.loads(chunk_json)
-                chunk = Chunk(**chunk_dict)
-            except Exception:
-                import traceback
-
-                traceback.print_exc()
-                print(f"Failed to parse document: {chunk_json}")
-                continue
-
-            chunks.append(chunk)
-            scores.append(1.0 / doc.metadata.distance)
-
-        return QueryDocumentsResponse(chunks=chunks, scores=scores)
-
-
-class WeaviateMemoryAdapter(
-    Memory, NeedsRequestProviderData, MemoryBanksProtocolPrivate
-):
-    def __init__(self, config: WeaviateConfig) -> None:
-        self.config = config
-        self.client_cache = {}
-        self.cache = {}
-
-    def _get_client(self) -> weaviate.Client:
-        provider_data = self.get_request_provider_data()
-        assert provider_data is not None, "Request provider data must be set"
-        assert isinstance(provider_data, WeaviateRequestProviderData)
-
-        key = f"{provider_data.weaviate_cluster_url}::{provider_data.weaviate_api_key}"
-        if key in self.client_cache:
-            return self.client_cache[key]
-
-        client = weaviate.connect_to_weaviate_cloud(
-            cluster_url=provider_data.weaviate_cluster_url,
-            auth_credentials=Auth.api_key(provider_data.weaviate_api_key),
-        )
-        self.client_cache[key] = client
-        return client
-
-    async def initialize(self) -> None:
-        pass
-
-    async def shutdown(self) -> None:
-        for client in self.client_cache.values():
-            client.close()
-
-    async def register_memory_bank(
-        self,
-        memory_bank: MemoryBankDef,
-    ) -> None:
-        assert (
-            memory_bank.type == MemoryBankType.vector.value
-        ), f"Only vector banks are supported {memory_bank.type}"
-
-        client = self._get_client()
-
-        # Create collection if it doesn't exist
-        if not client.collections.exists(memory_bank.identifier):
-            client.collections.create(
-                name=memory_bank.identifier,
-                vectorizer_config=wvc.config.Configure.Vectorizer.none(),
-                properties=[
-                    wvc.config.Property(
-                        name="chunk_content",
-                        data_type=wvc.config.DataType.TEXT,
-                    ),
-                ],
-            )
-
-        index = BankWithIndex(
-            bank=memory_bank,
-            index=WeaviateIndex(client=client, collection_name=memory_bank.identifier),
-        )
-        self.cache[memory_bank.identifier] = index
-
-    async def list_memory_banks(self) -> List[MemoryBankDef]:
-        # TODO: right now the Llama Stack is the source of truth for these banks. That is
-        # not ideal. It should be Weaviate which is the source of truth. Unfortunately,
-        # list() happens at Stack startup when the Weaviate client (credentials) is not
-        # yet available. We need to figure out a way to make this work.
-        return [i.bank for i in self.cache.values()]
-
-    async def _get_and_cache_bank_index(self, bank_id: str) -> Optional[BankWithIndex]:
-        if bank_id in self.cache:
-            return self.cache[bank_id]
-
-        bank = await self.memory_bank_store.get_memory_bank(bank_id)
-        if not bank:
-            raise ValueError(f"Bank {bank_id} not found")
-
-        client = self._get_client()
-        if not client.collections.exists(bank_id):
-            raise ValueError(f"Collection with name `{bank_id}` not found")
-
-        index = BankWithIndex(
-            bank=bank,
-            index=WeaviateIndex(client=client, collection_name=bank_id),
-        )
-        self.cache[bank_id] = index
-        return index
-
-    async def insert_documents(
-        self,
-        bank_id: str,
-        documents: List[MemoryBankDocument],
-        ttl_seconds: Optional[int] = None,
-    ) -> None:
-        index = await self._get_and_cache_bank_index(bank_id)
-        if not index:
-            raise ValueError(f"Bank {bank_id} not found")
-
-        await index.insert_documents(documents)
-
-    async def query_documents(
-        self,
-        bank_id: str,
-        query: InterleavedTextMedia,
-        params: Optional[Dict[str, Any]] = None,
-    ) -> QueryDocumentsResponse:
-        index = await self._get_and_cache_bank_index(bank_id)
-        if not index:
-            raise ValueError(f"Bank {bank_id} not found")
-
-        return await index.query_documents(query, params)

llama_stack/providers/adapters/safety/bedrock/bedrock.py

@@ -1,113 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-import json
-import logging
-
-from typing import Any, Dict, List
-
-import boto3
-
-from llama_stack.apis.safety import *  # noqa
-from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.providers.datatypes import ShieldsProtocolPrivate
-
-from .config import BedrockSafetyConfig
-
-
-logger = logging.getLogger(__name__)
-
-
-BEDROCK_SUPPORTED_SHIELDS = [
-    ShieldType.generic_content_shield.value,
-]
-
-
-class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
-    def __init__(self, config: BedrockSafetyConfig) -> None:
-        if not config.aws_profile:
-            raise ValueError(f"Missing boto_client aws_profile in model info::{config}")
-        self.config = config
-        self.registered_shields = []
-
-    async def initialize(self) -> None:
-        try:
-            print(f"initializing with profile --- > {self.config}")
-            self.boto_client = boto3.Session(
-                profile_name=self.config.aws_profile
-            ).client("bedrock-runtime")
-        except Exception as e:
-            raise RuntimeError("Error initializing BedrockSafetyAdapter") from e
-
-    async def shutdown(self) -> None:
-        pass
-
-    async def register_shield(self, shield: ShieldDef) -> None:
-        raise ValueError("Registering dynamic shields is not supported")
-
-    async def list_shields(self) -> List[ShieldDef]:
-        raise NotImplementedError(
-            """
-            `list_shields` not implemented; this should read all guardrails from
-            bedrock and populate guardrailId and guardrailVersion in the ShieldDef.
-        """
-        )
-
-    async def run_shield(
-        self, shield_type: str, messages: List[Message], params: Dict[str, Any] = None
-    ) -> RunShieldResponse:
-        shield_def = await self.shield_store.get_shield(shield_type)
-        if not shield_def:
-            raise ValueError(f"Unknown shield {shield_type}")
-
-        """This is the implementation for the bedrock guardrails. The input to the guardrails is to be of this format
-        ```content = [
-            {
-                "text": {
-                    "text": "Is the AB503 Product a better investment than the S&P 500?"
-                }
-            }
-        ]```
-        However the incoming messages are of this type UserMessage(content=....) coming from
-        https://github.com/meta-llama/llama-models/blob/main/models/llama3/api/datatypes.py
-
-        They contain content, role . For now we will extract the content and default the "qualifiers": ["query"]
-        """
-
-        shield_params = shield_def.params
-        logger.debug(f"run_shield::{shield_params}::messages={messages}")
-
-        # - convert the messages into format Bedrock expects
-        content_messages = []
-        for message in messages:
-            content_messages.append({"text": {"text": message.content}})
-        logger.debug(
-            f"run_shield::final:messages::{json.dumps(content_messages, indent=2)}:"
-        )
-
-        response = self.boto_client.apply_guardrail(
-            guardrailIdentifier=shield_params["guardrailIdentifier"],
-            guardrailVersion=shield_params["guardrailVersion"],
-            source="OUTPUT",  # or 'INPUT' depending on your use case
-            content=content_messages,
-        )
-        if response["action"] == "GUARDRAIL_INTERVENED":
-            user_message = ""
-            metadata = {}
-            for output in response["outputs"]:
-                # guardrails returns a list - however for this implementation we will leverage the last values
-                user_message = output["text"]
-            for assessment in response["assessments"]:
-                # guardrails returns a list - however for this implementation we will leverage the last values
-                metadata = dict(assessment)
-
-            return SafetyViolation(
-                user_message=user_message,
-                violation_level=ViolationLevel.ERROR,
-                metadata=metadata,
-            )
-
-        return None

llama_stack/providers/adapters/safety/bedrock/config.py

@@ -1,16 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from pydantic import BaseModel, Field
-
-
-class BedrockSafetyConfig(BaseModel):
-    """Configuration information for a guardrail that you want to use in the request."""
-
-    aws_profile: str = Field(
-        default="default",
-        description="The profile on the machine having valid aws credentials. This will ensure separation of creation to invocation",
-    )

llama_stack/providers/adapters/safety/sample/sample.py

@@ -1,23 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from .config import SampleConfig
-
-
-from llama_stack.apis.safety import *  # noqa: F403
-
-
-class SampleSafetyImpl(Safety):
-    def __init__(self, config: SampleConfig):
-        self.config = config
-
-    async def register_shield(self, shield: ShieldDef) -> None:
-        # these are the safety shields the Llama Stack will use to route requests to this provider
-        # perform validation here if necessary
-        pass
-
-    async def initialize(self):
-        pass

llama_stack/providers/adapters/safety/together/__init__.py

@@ -1,18 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from .config import TogetherProviderDataValidator, TogetherSafetyConfig  # noqa: F401
-
-
-async def get_adapter_impl(config: TogetherSafetyConfig, _deps):
-    from .together import TogetherSafetyImpl
-
-    assert isinstance(
-        config, TogetherSafetyConfig
-    ), f"Unexpected config type: {type(config)}"
-    impl = TogetherSafetyImpl(config)
-    await impl.initialize()
-    return impl

llama_stack/providers/adapters/safety/together/config.py

@@ -1,26 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from typing import Optional
-
-from llama_models.schema_utils import json_schema_type
-from pydantic import BaseModel, Field
-
-
-class TogetherProviderDataValidator(BaseModel):
-    together_api_key: str
-
-
-@json_schema_type
-class TogetherSafetyConfig(BaseModel):
-    url: str = Field(
-        default="https://api.together.xyz/v1",
-        description="The URL for the Together AI server",
-    )
-    api_key: Optional[str] = Field(
-        default=None,
-        description="The Together AI API Key (default for the distribution, if any)",
-    )