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

frogml_core/vector_store/utils/upsert_utils.py

@@ -1,218 +0,0 @@
-import os
-from functools import partial
-from math import ceil
-from multiprocessing import Pool, set_start_method
-from typing import Dict, Iterable, List, Optional, Tuple, TypeVar, Union
-
-from dependency_injector.wiring import Provide, inject
-from typeguard import typechecked
-
-from frogml_proto.qwak.vectors.v1.vector_pb2 import (
-    DoubleVector,
-    Property,
-    StoredVector,
-    VectorIdentifier,
-)
-from frogml_core.clients.vector_store.serving_client import VectorServingClient
-from frogml_core.exceptions import FrogmlException
-from frogml_core.inner.di_configuration import FrogmlContainer
-from frogml_core.vector_store.inference_client import VectorStoreInferenceClient
-
-_NaturalInput = TypeVar("T")
-_NaturalInputs = List[_NaturalInput]
-_Vector = List[float]
-_Properties = Dict[str, Union[str, int, bool, float]]
-
-_type_to_proto_property_mapping: Dict = {
-    str: "string_val",
-    bool: "bool_val",
-    int: "int_val",
-    float: "double_val",
-}
-
-
-def _build_property(key: str, value: Union[str, int, bool, float]):
-    type_val = _type_to_proto_property_mapping.get(type(value), None)
-    if not type_val:
-        raise FrogmlException(
-            f"Cannot upsert vector with property value type {type(value)}. "
-            f"Supported types are: {list(_type_to_proto_property_mapping.keys())}"
-        )
-
-    property_args = {"name": key, type_val: value}
-    return Property(**property_args)
-
-
-def _rewire_qwak_container(config):
-    # re-creating the container using the config from the original container.
-    # note that this runs in a fresh interpreter - at that point there's a running
-    # container because of the imports, but it won't necessarily have the same config
-    # as the container in the parent process that spawned this one.
-    # rewiring only the vector store + ecosystem and authentication - if using stuff from
-    # other modules pls feel free to add it here.
-    new_container = FrogmlContainer(config=config)
-    from frogml_core.clients import vector_store
-    from frogml_core.clients.administration import authentication, eco_system
-
-    new_container.wire(
-        packages=[
-            authentication,
-            eco_system,
-            vector_store,
-        ]
-    )
-
-
-@typechecked
-def _divide_chunks(lst: List, chunk_size: int):
-    if chunk_size <= 0:
-        raise FrogmlException("Chunk size must be a positive integer")
-
-    num_items: int = len(lst)
-    num_chunks: int = ceil(num_items / chunk_size)
-    for i in range(num_chunks):
-        yield lst[i * chunk_size : (i + 1) * chunk_size]
-
-
-def _get_vector_identifier(t: Union[str, Tuple[str, str]]):
-    if type(t) is str:
-        return VectorIdentifier(vector_id=t)
-    return VectorIdentifier(vector_id=t[0], tenant_id=t[1])
-
-
-def _upsert_vector_block(
-    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _Vector, _Properties]],
-    chunk_size: int,
-    collection_name: str,
-    edge_services_url: str,
-) -> None:
-    vector_serving_client: VectorServingClient = VectorServingClient(
-        edge_services_url=edge_services_url
-    )
-    for chunk in _divide_chunks(vector_tuples, chunk_size):
-        # chunk is a list of (id, vector, properties) tuples
-        vector_serving_client.upsert_vectors(
-            collection_name=collection_name,
-            vectors=[
-                StoredVector(
-                    vector_identifier=_get_vector_identifier(tpl[0]),
-                    vector=DoubleVector(element=tpl[1]),
-                    properties=[
-                        _build_property(key, value) for (key, value) in tpl[2].items()
-                    ],
-                )
-                for tpl in chunk
-            ],
-        )
-
-
-def _upsert_natural_input_block(
-    vector_tuples: List[Tuple[VectorIdentifier, _NaturalInput, _Properties]],
-    chunk_size: int,
-    vectorizer_name: str,
-    collection_name: str,
-    edge_services_url: str,
-) -> None:
-    vector_serving_client: VectorServingClient = VectorServingClient(
-        edge_services_url=edge_services_url
-    )
-    inference_client: VectorStoreInferenceClient = VectorStoreInferenceClient(
-        model_id=vectorizer_name
-    )
-    for chunk in _divide_chunks(vector_tuples, chunk_size):
-        # chunk is a list of (id, _NaturalInput, properties) tuples
-        vector_serving_client.upsert_vectors(
-            collection_name=collection_name,
-            vectors=[
-                StoredVector(
-                    vector_identifier=_get_vector_identifier(tpl[0]),
-                    vector=DoubleVector(
-                        element=inference_client.get_embedding(natural_input=tpl[1])
-                    ),
-                    properties=[
-                        _build_property(key, value) for (key, value) in tpl[2].items()
-                    ],
-                )
-                for tpl in chunk
-            ],
-        )
-
-
-@inject
-def _upsert_natural_input(
-    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _NaturalInput, _Properties]],
-    chunk_size: int,
-    vectorizer_name: str,
-    collection_name: str,
-    edge_services_url: str,
-    multiproc: bool = False,
-    max_processes: Optional[int] = None,
-    config=Provide[FrogmlContainer.config],
-):
-    if not multiproc:
-        _upsert_natural_input_block(
-            vector_tuples=vector_tuples,
-            chunk_size=chunk_size,
-            vectorizer_name=vectorizer_name,
-            collection_name=collection_name,
-            edge_services_url=edge_services_url,
-        )
-    else:
-        if max_processes is None:
-            max_processes = os.cpu_count()
-        effective_block_size: int = ceil(len(vector_tuples) / (max_processes * 4))
-
-        f = partial(
-            _upsert_natural_input_block,
-            chunk_size=chunk_size,
-            vectorizer_name=vectorizer_name,
-            collection_name=collection_name,
-            edge_services_url=edge_services_url,
-        )
-
-        blocks: Iterable[List[Tuple[str, _NaturalInput, _Properties]]] = _divide_chunks(
-            vector_tuples, effective_block_size
-        )
-        initializer = partial(_rewire_qwak_container, config=config)
-        set_start_method("spawn", force=True)
-
-        with Pool(processes=max_processes, initializer=initializer) as p:
-            p.map(f, blocks)
-
-
-@inject
-def _upsert_vectors(
-    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _Vector, _Properties]],
-    chunk_size: int,
-    collection_name: str,
-    edge_services_url: str,
-    multiproc: bool = False,
-    max_processes: Optional[int] = None,
-    config=Provide[FrogmlContainer.config],
-):
-    if not multiproc:
-        _upsert_vector_block(
-            vector_tuples=vector_tuples,
-            chunk_size=chunk_size,
-            collection_name=collection_name,
-            edge_services_url=edge_services_url,
-        )
-    else:
-        if max_processes is None:
-            max_processes = os.cpu_count()
-        effective_block_size: int = ceil(len(vector_tuples) / (max_processes * 4))
-
-        f = partial(
-            _upsert_vector_block,
-            chunk_size=chunk_size,
-            collection_name=collection_name,
-            edge_services_url=edge_services_url,
-        )
-        blocks: Iterable[List[Tuple[str, _Vector, _Properties]]] = list(
-            _divide_chunks(vector_tuples, effective_block_size)
-        )
-
-        set_start_method("spawn", force=True)
-        initializer = partial(_rewire_qwak_container, config=config)
-        with Pool(processes=max_processes, initializer=initializer) as p:
-            p.map(f, blocks)

frogml_proto/qwak/prompt/v1/prompt/prompt_manager_service_pb2.py

@@ -1,77 +0,0 @@
-# -*- coding: utf-8 -*-
-# Generated by the protocol buffer compiler.  DO NOT EDIT!
-# source: frogml_proto.qwak.prompt/v1/prompt/prompt_manager_service.proto
-"""Generated protocol buffer code."""
-from google.protobuf import descriptor as _descriptor
-from google.protobuf import descriptor_pool as _descriptor_pool
-from google.protobuf import symbol_database as _symbol_database
-from google.protobuf.internal import builder as _builder
-# @@protoc_insertion_point(imports)
-
-_sym_db = _symbol_database.Default()
-
-
-from frogml_proto.qwak.prompt.v1.prompt import prompt_pb2 as qwak_dot_prompt_dot_v1_dot_prompt_dot_prompt__pb2
-from frogml_proto.qwak.model_descriptor import open_ai_descriptor_pb2 as qwak_dot_model__descriptor_dot_open__ai__descriptor__pb2
-
-
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n2qwak/prompt/v1/prompt/prompt_manager_service.proto\x12\x15qwak.prompt.v1.prompt\x1a\"qwak/prompt/v1/prompt/prompt.proto\x1a.qwak/model_descriptor/open_ai_descriptor.proto\"\xa9\x01\n\x13\x43reatePromptRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\x12\x45\n\x13prompt_version_spec\x18\x02 \x01(\x0b\x32(.qwak.prompt.v1.prompt.PromptVersionSpec\x12\x36\n\x0bprompt_spec\x18\x03 \x01(\x0b\x32!.qwak.prompt.v1.prompt.PromptSpec\"E\n\x14\x43reatePromptResponse\x12-\n\x06prompt\x18\x01 \x01(\x0b\x32\x1d.qwak.prompt.v1.prompt.Prompt\"\x8d\x01\n\x1a\x43reatePromptVersionRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\x12\x45\n\x13prompt_version_spec\x18\x02 \x01(\x0b\x32(.qwak.prompt.v1.prompt.PromptVersionSpec\x12\x13\n\x0bset_default\x18\x03 \x01(\x08\"[\n\x1b\x43reatePromptVersionResponse\x12<\n\x0eprompt_version\x18\x01 \x01(\x0b\x32$.qwak.prompt.v1.prompt.PromptVersion\"b\n\x13UpdatePromptRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\x12\x36\n\x0bprompt_spec\x18\x02 \x01(\x0b\x32!.qwak.prompt.v1.prompt.PromptSpec\"\x16\n\x14UpdatePromptResponse\"o\n#GetPromptVersionByPromptNameRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\x12\x18\n\x0eversion_number\x18\x02 \x01(\x05H\x00\x42\x19\n\x17optional_version_number\"d\n$GetPromptVersionByPromptNameResponse\x12<\n\x0eprompt_version\x18\x01 \x01(\x0b\x32$.qwak.prompt.v1.prompt.PromptVersion\"-\n\x16GetPromptByNameRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\"H\n\x17GetPromptByNameResponse\x12-\n\x06prompt\x18\x01 \x01(\x0b\x32\x1d.qwak.prompt.v1.prompt.Prompt\"\x14\n\x12ListPromptsRequest\"E\n\x13ListPromptsResponse\x12.\n\x07prompts\x18\x01 \x03(\x0b\x32\x1d.qwak.prompt.v1.prompt.Prompt\"8\n!InitPaginationByPromptNameRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\"V\n\"InitPaginationByPromptNameResponse\x12\x14\n\x0crecord_count\x18\x01 \x01(\x03\x12\x1a\n\x12max_version_number\x18\x02 \x01(\x05\"\x80\x01\n%ListPromptVersionsByPromptNameRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\x12\x11\n\tpage_size\x18\x02 \x01(\x05\x12\x13\n\x0bpage_number\x18\x03 \x01(\x05\x12\x1a\n\x12max_version_number\x18\x04 \x01(\x05\"g\n&ListPromptVersionsByPromptNameResponse\x12=\n\x0fprompt_versions\x18\x01 \x03(\x0b\x32$.qwak.prompt.v1.prompt.PromptVersion\"*\n\x13\x44\x65letePromptRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\"\x16\n\x14\x44\x65letePromptResponse\"I\n\x1a\x44\x65letePromptVersionRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\x12\x16\n\x0eversion_number\x18\x03 \x01(\x05\"\x1d\n\x1b\x44\x65letePromptVersionResponse\"M\n\x1eSetDefaultPromptVersionRequest\x12\x13\n\x0bprompt_name\x18\x01 \x01(\t\x12\x16\n\x0eversion_number\x18\x02 \x01(\x05\"!\n\x1fSetDefaultPromptVersionResponse\"\x19\n\x17ListPromptBriefsRequest\"U\n\x18ListPromptBriefsResponse\x12\x39\n\rprompt_briefs\x18\x01 \x03(\x0b\x32\".qwak.prompt.v1.prompt.PromptBrief2\xf3\x0b\n\x14PromptManagerService\x12g\n\x0c\x43reatePrompt\x12*.qwak.prompt.v1.prompt.CreatePromptRequest\x1a+.qwak.prompt.v1.prompt.CreatePromptResponse\x12|\n\x13\x43reatePromptVersion\x12\x31.qwak.prompt.v1.prompt.CreatePromptVersionRequest\x1a\x32.qwak.prompt.v1.prompt.CreatePromptVersionResponse\x12g\n\x0cUpdatePrompt\x12*.qwak.prompt.v1.prompt.UpdatePromptRequest\x1a+.qwak.prompt.v1.prompt.UpdatePromptResponse\x12p\n\x0fGetPromptByName\x12-.qwak.prompt.v1.prompt.GetPromptByNameRequest\x1a..qwak.prompt.v1.prompt.GetPromptByNameResponse\x12\x64\n\x0bListPrompts\x12).qwak.prompt.v1.prompt.ListPromptsRequest\x1a*.qwak.prompt.v1.prompt.ListPromptsResponse\x12\x97\x01\n\x1cGetPromptVersionByPromptName\x12:.qwak.prompt.v1.prompt.GetPromptVersionByPromptNameRequest\x1a;.qwak.prompt.v1.prompt.GetPromptVersionByPromptNameResponse\x12\x91\x01\n\x1aInitPaginationByPromptName\x12\x38.qwak.prompt.v1.prompt.InitPaginationByPromptNameRequest\x1a\x39.qwak.prompt.v1.prompt.InitPaginationByPromptNameResponse\x12\x9d\x01\n\x1eListPromptVersionsByPromptName\x12<.qwak.prompt.v1.prompt.ListPromptVersionsByPromptNameRequest\x1a=.qwak.prompt.v1.prompt.ListPromptVersionsByPromptNameResponse\x12g\n\x0c\x44\x65letePrompt\x12*.qwak.prompt.v1.prompt.DeletePromptRequest\x1a+.qwak.prompt.v1.prompt.DeletePromptResponse\x12|\n\x13\x44\x65letePromptVersion\x12\x31.qwak.prompt.v1.prompt.DeletePromptVersionRequest\x1a\x32.qwak.prompt.v1.prompt.DeletePromptVersionResponse\x12\x88\x01\n\x17SetDefaultPromptVersion\x12\x35.qwak.prompt.v1.prompt.SetDefaultPromptVersionRequest\x1a\x36.qwak.prompt.v1.prompt.SetDefaultPromptVersionResponse\x12s\n\x10ListPromptBriefs\x12..qwak.prompt.v1.prompt.ListPromptBriefsRequest\x1a/.qwak.prompt.v1.prompt.ListPromptBriefsResponseB$\n com.qwak.ai.prompt.api.v1.promptP\x01\x62\x06proto3')
-
-_globals = globals()
-_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
-_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'qwak.prompt.v1.prompt.prompt_manager_service_pb2', _globals)
-if _descriptor._USE_C_DESCRIPTORS == False:
-
-  DESCRIPTOR._options = None
-  DESCRIPTOR._serialized_options = b'\n com.qwak.ai.prompt.api.v1.promptP\001'
-  _globals['_CREATEPROMPTREQUEST']._serialized_start=162
-  _globals['_CREATEPROMPTREQUEST']._serialized_end=331
-  _globals['_CREATEPROMPTRESPONSE']._serialized_start=333
-  _globals['_CREATEPROMPTRESPONSE']._serialized_end=402
-  _globals['_CREATEPROMPTVERSIONREQUEST']._serialized_start=405
-  _globals['_CREATEPROMPTVERSIONREQUEST']._serialized_end=546
-  _globals['_CREATEPROMPTVERSIONRESPONSE']._serialized_start=548
-  _globals['_CREATEPROMPTVERSIONRESPONSE']._serialized_end=639
-  _globals['_UPDATEPROMPTREQUEST']._serialized_start=641
-  _globals['_UPDATEPROMPTREQUEST']._serialized_end=739
-  _globals['_UPDATEPROMPTRESPONSE']._serialized_start=741
-  _globals['_UPDATEPROMPTRESPONSE']._serialized_end=763
-  _globals['_GETPROMPTVERSIONBYPROMPTNAMEREQUEST']._serialized_start=765
-  _globals['_GETPROMPTVERSIONBYPROMPTNAMEREQUEST']._serialized_end=876
-  _globals['_GETPROMPTVERSIONBYPROMPTNAMERESPONSE']._serialized_start=878
-  _globals['_GETPROMPTVERSIONBYPROMPTNAMERESPONSE']._serialized_end=978
-  _globals['_GETPROMPTBYNAMEREQUEST']._serialized_start=980
-  _globals['_GETPROMPTBYNAMEREQUEST']._serialized_end=1025
-  _globals['_GETPROMPTBYNAMERESPONSE']._serialized_start=1027
-  _globals['_GETPROMPTBYNAMERESPONSE']._serialized_end=1099
-  _globals['_LISTPROMPTSREQUEST']._serialized_start=1101
-  _globals['_LISTPROMPTSREQUEST']._serialized_end=1121
-  _globals['_LISTPROMPTSRESPONSE']._serialized_start=1123
-  _globals['_LISTPROMPTSRESPONSE']._serialized_end=1192
-  _globals['_INITPAGINATIONBYPROMPTNAMEREQUEST']._serialized_start=1194
-  _globals['_INITPAGINATIONBYPROMPTNAMEREQUEST']._serialized_end=1250
-  _globals['_INITPAGINATIONBYPROMPTNAMERESPONSE']._serialized_start=1252
-  _globals['_INITPAGINATIONBYPROMPTNAMERESPONSE']._serialized_end=1338
-  _globals['_LISTPROMPTVERSIONSBYPROMPTNAMEREQUEST']._serialized_start=1341
-  _globals['_LISTPROMPTVERSIONSBYPROMPTNAMEREQUEST']._serialized_end=1469
-  _globals['_LISTPROMPTVERSIONSBYPROMPTNAMERESPONSE']._serialized_start=1471
-  _globals['_LISTPROMPTVERSIONSBYPROMPTNAMERESPONSE']._serialized_end=1574
-  _globals['_DELETEPROMPTREQUEST']._serialized_start=1576
-  _globals['_DELETEPROMPTREQUEST']._serialized_end=1618
-  _globals['_DELETEPROMPTRESPONSE']._serialized_start=1620
-  _globals['_DELETEPROMPTRESPONSE']._serialized_end=1642
-  _globals['_DELETEPROMPTVERSIONREQUEST']._serialized_start=1644
-  _globals['_DELETEPROMPTVERSIONREQUEST']._serialized_end=1717
-  _globals['_DELETEPROMPTVERSIONRESPONSE']._serialized_start=1719
-  _globals['_DELETEPROMPTVERSIONRESPONSE']._serialized_end=1748
-  _globals['_SETDEFAULTPROMPTVERSIONREQUEST']._serialized_start=1750
-  _globals['_SETDEFAULTPROMPTVERSIONREQUEST']._serialized_end=1827
-  _globals['_SETDEFAULTPROMPTVERSIONRESPONSE']._serialized_start=1829
-  _globals['_SETDEFAULTPROMPTVERSIONRESPONSE']._serialized_end=1862
-  _globals['_LISTPROMPTBRIEFSREQUEST']._serialized_start=1864
-  _globals['_LISTPROMPTBRIEFSREQUEST']._serialized_end=1889
-  _globals['_LISTPROMPTBRIEFSRESPONSE']._serialized_start=1891
-  _globals['_LISTPROMPTBRIEFSRESPONSE']._serialized_end=1976
-  _globals['_PROMPTMANAGERSERVICE']._serialized_start=1979
-  _globals['_PROMPTMANAGERSERVICE']._serialized_end=3502
-# @@protoc_insertion_point(module_scope)