frogml-core 0.0.113__py3-none-any.whl → 0.0.114__py3-none-any.whl

frogml_services_mock/mocks/prompt_manager_service.py

@@ -1,281 +0,0 @@
-from typing import Dict, List, Tuple
-
-import grpc
-from google.protobuf.timestamp_pb2 import Timestamp
-
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    CreatePromptRequest as ProtoCreatePromptRequest,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    CreatePromptResponse as ProtoCreatePromptResponse,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    CreatePromptVersionRequest as ProtoCreatePromptVersionRequest,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    CreatePromptVersionResponse as ProtoCreatePromptVersionResponse,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    DeletePromptRequest as ProtoDeletePromptRequest,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    DeletePromptResponse as ProtoDeletePromptResponse,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    DeletePromptVersionRequest as ProtoDeletePromptVersionRequest,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    DeletePromptVersionResponse as ProtoDeletePromptVersionResponse,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    GetPromptByNameRequest as ProtoGetPromptByNameRequest,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    GetPromptByNameResponse as ProtoGetPromptByNameResponse,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    GetPromptVersionByPromptNameRequest as ProtoGetPromptVersionByPromptNameRequest,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    GetPromptVersionByPromptNameResponse as ProtoGetPromptVersionByPromptNameResponse,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    SetDefaultPromptVersionRequest as ProtoSetDefaultPromptVersionRequest,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    SetDefaultPromptVersionResponse as ProtoSetDefaultPromptVersionResponse,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    UpdatePromptRequest as ProtoUpdatePromptRequest,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2 import (
-    UpdatePromptResponse as ProtoUpdatePromptResponse,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_manager_service_pb2_grpc import (
-    PromptManagerServiceServicer,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import Prompt as ProtoPrompt
-from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
-    PromptMetadata as ProtoPromptMetadata,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import PromptSpec as ProtoPromptSpec
-from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
-    PromptVersion as ProtoPromptVersion,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
-    PromptVersionDefinition as ProtoPromptVersionDefinition,
-)
-from frogml_proto.qwak.prompt.v1.prompt.prompt_pb2 import (
-    PromptVersionMetadata as ProtoPromptVersionMetadata,
-)
-
-
-class PromptManagerServiceMock(PromptManagerServiceServicer):
-    prompts: Dict[str, ProtoPrompt] = {}
-    prompts_versions: Dict[Tuple[str, int], ProtoPromptVersionDefinition] = {}
-    prompts_next_version: Dict[str, int] = {}
-
-    def clear(self):
-        self.prompts = {}
-        self.prompts_versions = {}
-        self.prompts_next_version = {}
-
-    def _is_prompt_name_valid(self, test_str: str) -> bool:
-        import re
-
-        regex = r"^[a-z0-9](?:[-_]?[a-z0-9]+)+$"
-        return re.match(pattern=regex, string=test_str) is not None
-
-    def _raise_on_missing_prompt(self, prompt_name: str, context):
-        if prompt_name not in self.prompts:
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details("Prompt name does not exist")
-            raise ValueError("Prompt name does not exist")
-
-    def _raise_on_missing_prompt_version(self, prompt_name: str, version: int, context):
-        if (
-            prompt_name not in self.prompts
-            or (prompt_name, version) not in self.prompts_versions
-        ):
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details("Prompt version does not exist")
-            raise ValueError("Prompt version does not exist")
-
-    def CreatePrompt(self, request: ProtoCreatePromptRequest, context):
-        """Create a new prompt"""
-
-        if request.prompt_name in self.prompts:
-            context.set_code(grpc.StatusCode.ALREADY_EXISTS)
-            context.set_details("Prompt name already exists")
-            raise ValueError("Prompt name already exists")
-
-        if not self._is_prompt_name_valid(request.prompt_name):
-            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
-            context.set_details("Bad Prompt name")
-            raise ValueError("Bad Prompt name")
-
-        prompt_version_definition = ProtoPromptVersionDefinition(
-            version_number=1,
-            version_spec=request.prompt_version_spec,
-            version_metadata=ProtoPromptVersionMetadata(
-                created_by="nirh@frogml.com", created_at=Timestamp().GetCurrentTime()
-            ),
-        )
-
-        prompt = ProtoPrompt(
-            name=request.prompt_name,
-            default_version_definition=prompt_version_definition,
-            metadata=ProtoPromptMetadata(
-                created_at=Timestamp().GetCurrentTime(), created_by="nirh@frogml.com"
-            ),
-            prompt_spec=ProtoPromptSpec(description=request.prompt_spec.description),
-        )
-
-        self.prompts[prompt.name] = prompt
-        self.prompts_next_version[prompt.name] = 2
-        self.prompts_versions[
-            (prompt.name, prompt.default_version_definition.version_number)
-        ] = prompt_version_definition
-
-        return ProtoCreatePromptResponse(prompt=prompt)
-
-    def CreatePromptVersion(self, request: ProtoCreatePromptVersionRequest, context):
-        """
-        Create new prompt version.
-        Set new version in versions dict.
-        Bump prompt next version.
-        Set default version if needed
-        """
-        self._raise_on_missing_prompt(prompt_name=request.prompt_name, context=context)
-
-        existing_prompt: ProtoPrompt = self.prompts[request.prompt_name]
-        new_prompt_version_number: int = self.prompts_next_version[request.prompt_name]
-        prompt_version_definition = ProtoPromptVersionDefinition(
-            version_number=new_prompt_version_number,
-            version_spec=request.prompt_version_spec,
-            version_metadata=ProtoPromptVersionMetadata(
-                created_by="nirh@frogml.com", created_at=Timestamp().GetCurrentTime()
-            ),
-        )
-
-        self.prompts_next_version[request.prompt_name] = new_prompt_version_number + 1
-
-        if request.set_default:
-            existing_prompt.default_version_definition.CopyFrom(
-                prompt_version_definition
-            )
-
-        self.prompts_versions[(request.prompt_name, new_prompt_version_number)] = (
-            prompt_version_definition
-        )
-        return ProtoCreatePromptVersionResponse(
-            prompt_version=ProtoPromptVersion(
-                prompt_name=request.prompt_name,
-                prompt_version_definition=prompt_version_definition,
-            )
-        )
-
-    def UpdatePrompt(self, request: ProtoUpdatePromptRequest, context):
-        """Update Prompt"""
-        self._raise_on_missing_prompt(prompt_name=request.prompt_name, context=context)
-
-        existing_prompt: ProtoPrompt = self.prompts[request.prompt_name]
-        existing_prompt.prompt_spec.CopyFrom(request.prompt_spec)
-
-        return ProtoUpdatePromptResponse()
-
-    def GetPromptByName(self, request: ProtoGetPromptByNameRequest, context):
-        """Get prompt with its default version by name"""
-        self._raise_on_missing_prompt(prompt_name=request.prompt_name, context=context)
-
-        existing_prompt: ProtoPrompt = self.prompts[request.prompt_name]
-        return ProtoGetPromptByNameResponse(prompt=existing_prompt)
-
-    def GetPromptVersionByPromptName(
-        self, request: ProtoGetPromptVersionByPromptNameRequest, context
-    ):
-        """Get prompt version by name and optional version number"""
-
-        self._raise_on_missing_prompt(
-            prompt_name=request.prompt_name,
-            context=context,
-        )
-
-        self._raise_on_missing_prompt_version(
-            prompt_name=request.prompt_name,
-            version=request.version_number,
-            context=context,
-        )
-
-        return ProtoGetPromptVersionByPromptNameResponse(
-            prompt_version=ProtoPromptVersion(
-                prompt_name=request.prompt_name,
-                prompt_version_definition=self.prompts_versions[
-                    (request.prompt_name, request.version_number)
-                ],
-            )
-        )
-
-    def DeletePrompt(self, request: ProtoDeletePromptRequest, context):
-        """Delete prompt"""
-        self._raise_on_missing_prompt(prompt_name=request.prompt_name, context=context)
-
-        del self.prompts[request.prompt_name]
-        del self.prompts_next_version[request.prompt_name]
-
-        versions_to_delete: List[int] = [
-            version
-            for (name, version) in self.prompts_versions.keys()
-            if name == request.prompt_name
-        ]
-
-        for v in versions_to_delete:
-            del self.prompts_versions[(request.prompt_name, v)]
-
-        return ProtoDeletePromptResponse()
-
-    def DeletePromptVersion(self, request: ProtoDeletePromptVersionRequest, context):
-        """Delete prompt version"""
-        self._raise_on_missing_prompt(
-            prompt_name=request.prompt_name,
-            context=context,
-        )
-
-        self._raise_on_missing_prompt_version(
-            prompt_name=request.prompt_name,
-            version=request.version_number,
-            context=context,
-        )
-
-        prompt_default_version: ProtoPrompt = self.prompts[request.prompt_name]
-        if (
-            prompt_default_version.default_version_definition.version_number
-            == request.version_number
-        ):
-            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
-            context.set_details("Can't delete prompt default version")
-            raise ValueError("Can't delete prompt default version")
-
-        del self.prompts_versions[(request.prompt_name, request.version_number)]
-        return ProtoDeletePromptVersionResponse()
-
-    def SetDefaultPromptVersion(
-        self, request: ProtoSetDefaultPromptVersionRequest, context
-    ):
-        """Set default prompt version"""
-        self._raise_on_missing_prompt(
-            prompt_name=request.prompt_name,
-            context=context,
-        )
-
-        self._raise_on_missing_prompt_version(
-            prompt_name=request.prompt_name,
-            version=request.version_number,
-            context=context,
-        )
-        existing_prompt: ProtoPrompt = self.prompts[request.prompt_name]
-        prompt_version_definition: ProtoPromptVersionDefinition = self.prompts_versions[
-            (request.prompt_name, request.version_number)
-        ]
-        existing_prompt.default_version_definition.CopyFrom(prompt_version_definition)
-        return ProtoSetDefaultPromptVersionResponse()

frogml_services_mock/mocks/vector_serving_api.py

@@ -1,159 +0,0 @@
-from dataclasses import dataclass
-from typing import Dict, List
-
-import grpc
-from numpy import dot
-from numpy.linalg import norm
-
-from frogml_proto.qwak.vectors.v1.vector_pb2 import (
-    DoubleVector,
-    Property,
-    SearchResult,
-    StoredVector,
-    VectorIdentifier,
-)
-from frogml_proto.qwak.vectors.v1.vector_service_pb2 import (
-    DeleteVectorsResponse,
-    FetchVectorRequest,
-    FetchVectorResponse,
-    SearchSimilarVectorsRequest,
-    SearchSimilarVectorsResponse,
-    UpsertVectorsRequest,
-    UpsertVectorsResponse,
-)
-from frogml_proto.qwak.vectors.v1.vector_service_pb2_grpc import VectorServiceServicer
-
-
-@dataclass
-class VectorObject:
-    id: str
-    vector: List[float]
-    property: List[Property]
-
-
-class VectorServingServiceMock(VectorServiceServicer):
-    def __init__(self):
-        self._vector_collections: Dict[str, Dict] = dict()
-
-    def reset_vector_store(self) -> None:
-        self._vector_collections.clear()
-
-    def get_num_of_vectors(self, collection_name: str) -> int:
-        if collection_name not in self._vector_collections:
-            return 0
-
-        return len(self._vector_collections[collection_name].values())
-
-    def create_collection(self, collection_name: str) -> None:
-        self._vector_collections[collection_name] = dict()
-
-    def get_vector_by_ids(
-        self, collection_name: str, vector_ids: List[str]
-    ) -> List[dict]:
-        if collection_name not in self._vector_collections:
-            raise ValueError(f"Collection named {collection_name} doesn't exist")
-
-        collection = self._vector_collections[collection_name]
-        return [collection.get(_id) for _id in vector_ids if _id in collection]
-
-    def SearchSimilarVectors(self, request: SearchSimilarVectorsRequest, context):
-        if request.collection_name not in self._vector_collections:
-            context.set_details(f"Collection named {request.collection} doesn't exist'")
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            return
-
-        reference_vector = list(request.reference_vector.element)
-        collection_store = self._vector_collections.get(request.collection_name, dict())
-
-        # a naive impl of a "vector similarity" search - compute pairwise cosine distance on
-        # the entire set and return top results
-        result_set = sorted(
-            collection_store.values(),
-            key=lambda b: _cos_sim(reference_vector, b["vector"]),
-        )[-request.max_results :]
-
-        return SearchSimilarVectorsResponse(
-            search_results=[
-                SearchResult(
-                    id=result["id"] if request.include_id else None,
-                    properties=[
-                        p
-                        for p in result_set[0]["properties"]
-                        if p.name in request.properties
-                    ],
-                    vector=(
-                        DoubleVector(element=result["vector"])
-                        if request.include_vector
-                        else None
-                    ),
-                    distance=(
-                        _cos_sim(reference_vector, result["vector"])
-                        if request.include_distance
-                        else None
-                    ),
-                )
-                for result in result_set
-            ]
-        )
-
-    def UpsertVectors(self, request: UpsertVectorsRequest, context):
-        collection_store = self._vector_collections.get(request.collection_name, dict())
-        for stored_vector in request.vector:
-            id = stored_vector.vector_identifier.vector_id
-            collection_store[id] = {
-                "id": id,
-                "vector": list(stored_vector.vector.element),
-                "properties": stored_vector.properties,
-            }
-
-        self._vector_collections[request.collection_name] = collection_store
-        return UpsertVectorsResponse()
-
-    def DeleteVectors(self, request, context):
-        if request.collection_name not in self._vector_collections:
-            context.set_details(
-                f"Collection named {request.collection_name} doesn't exist'"
-            )
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            return
-
-        collection_store = self._vector_collections[request.collection_name]
-        ids_in_collection = [
-            vector_identifier.vector_id
-            for vector_identifier in request.vector_identifiers
-            if vector_identifier.vector_id in collection_store
-        ]
-        for id in ids_in_collection:
-            collection_store.pop(id)
-
-        self._vector_collections[request.collection_name] = collection_store
-        return DeleteVectorsResponse(num_vectors_deleted=len(ids_in_collection))
-
-    def FetchVector(self, request, context):
-        if request.collection_name not in self._vector_collections:
-            context.set_details(
-                f"Collection named {request.collection_name} doesn't exist'"
-            )
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            return
-
-        collection_store = self._vector_collections[request.collection_name]
-        vector_identifier = request.vector_identifier
-        vector_id = vector_identifier.vector_id
-        if vector_id not in collection_store:
-            context.set_details(f"Vector with id {vector_id} doesn't exist'")
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            return
-
-        stored_vector_dict = collection_store[vector_id]
-        return FetchVectorResponse(
-            vector=StoredVector(
-                vector_identifier=request.vector_identifier,
-                vector=DoubleVector(element=stored_vector_dict["vector"]),
-                properties=stored_vector_dict["properties"],
-            )
-        )
-
-
-def _cos_sim(a, b):
-    return dot(a, b) / (norm(a) * norm(b))

frogml_services_mock/mocks/vectors_management_api.py

@@ -1,97 +0,0 @@
-import uuid
-from datetime import datetime
-
-import grpc
-from google.protobuf.timestamp_pb2 import Timestamp
-
-from frogml_proto.qwak.vectors.v1.collection.collection_pb2 import (
-    COLLECTION_STATUS_CREATE_REQUESTED,
-    QwakMetadata,
-    VectorCollection,
-    VectorCollectionDefinition,
-)
-from frogml_proto.qwak.vectors.v1.collection.collection_service_pb2 import (
-    CreateCollectionResponse,
-    DeleteCollectionByIdResponse,
-    DeleteCollectionByNameResponse,
-    GetCollectionByIdResponse,
-    GetCollectionByNameResponse,
-    ListCollectionsResponse,
-)
-from frogml_proto.qwak.vectors.v1.collection.collection_service_pb2_grpc import (
-    VectorCollectionServiceServicer,
-)
-
-
-class VectorCollectionManagementServiceMock(VectorCollectionServiceServicer):
-    def __init__(self):
-        self._collections_spec_by_ids = {}
-        self._collections_spec_by_name = {}
-
-    def reset_collections(self):
-        self._collections_spec_by_ids.clear()
-        self._collections_spec_by_name.clear()
-
-    def CreateCollection(self, request, context):
-        timestamp = Timestamp()
-        timestamp.FromDatetime(datetime.now())
-
-        collection_id = str(uuid.uuid4())
-        vector_collection = VectorCollection(
-            metadata=QwakMetadata(
-                created_at=timestamp,
-                created_by="it@frogml.com",
-                last_modified_at=timestamp,
-                last_modified_by="it@frogml.com",
-            ),
-            definition=VectorCollectionDefinition(
-                id=collection_id,
-                collection_spec=request.collection_spec,
-            ),
-            status=COLLECTION_STATUS_CREATE_REQUESTED,
-        )
-
-        self._collections_spec_by_ids[collection_id] = vector_collection
-        self._collections_spec_by_name[request.collection_spec.name] = vector_collection
-        return CreateCollectionResponse(vector_collection=vector_collection)
-
-    def GetCollectionById(self, request, context):
-        if request.id in self._collections_spec_by_ids:
-            return GetCollectionByIdResponse(
-                vector_collection=self._collections_spec_by_ids[request.id]
-            )
-
-        context.set_details(f"Collection ID {request.id} doesn't exist'")
-        context.set_code(grpc.StatusCode.NOT_FOUND)
-
-    def GetCollectionByName(self, request, context):
-        if request.name in self._collections_spec_by_name:
-            return GetCollectionByNameResponse(
-                vector_collection=self._collections_spec_by_name[request.name]
-            )
-
-        context.set_details(f"Collection name {request.name} doesn't exist'")
-        context.set_code(grpc.StatusCode.NOT_FOUND)
-
-    def DeleteCollectionById(self, request, context):
-        if request.id in self._collections_spec_by_ids:
-            self._collections_spec_by_ids.pop(request.id)
-            return DeleteCollectionByIdResponse()
-
-        context.set_details(f"Collection ID {request.id} doesn't exist'")
-        context.set_code(grpc.StatusCode.NOT_FOUND)
-
-    def DeleteCollectionByName(self, request, context):
-        if request.name in self._collections_spec_by_name:
-            self._collections_spec_by_name.pop(request.name)
-            return DeleteCollectionByNameResponse()
-
-        context.set_details(f"Collection name {request.name} doesn't exist'")
-        context.set_code(grpc.StatusCode.NOT_FOUND)
-
-    def ListCollections(self, request, context):
-        return ListCollectionsResponse(
-            vector_collections=[
-                collection for collection in self._collections_spec_by_ids.values()
-            ]
-        )