kQj0ddlZddlmZmZmZmZmZddlmZddl m Z m Z m Z GddZ eZded efd Ze Zd Zd e fd Zy) N)DictOptionalSetTupleType) TypeAdapter)EmbeddingFunctionEmbeddingFunctionConfigUntrustedEmbeddingFunctionceZdZdZedZdZddeefdZ dZ ded e e fd Z d eeeefd eed ffd Zdeded e fdZddedeed dfdZded dfdZdedeed efdZde d dfdZddZdZdeded dfdZded efdZy)EmbeddingFunctionRegistryaR This is a singleton class used to register embedding functions and fetch them by name. It also handles serializing and deserializing. You can implement your own embedding function by subclassing EmbeddingFunction or TextEmbeddingFunction and registering it with the registry. NOTE: Here TEXT is a type alias for Union[str, List[str], pa.Array, pa.ChunkedArray, np.ndarray] Examples -------- >>> registry = EmbeddingFunctionRegistry.get_instance() >>> @registry.register("my-embedding-function") ... class MyEmbeddingFunction(EmbeddingFunction): ... def ndims(self) -> int: ... return 128 ... ... def compute_query_embeddings(self, query: str, *args, **kwargs): ... return self.compute_source_embeddings(query, *args, **kwargs) ... ... def compute_source_embeddings(self, texts, *args, **kwargs): ... return [np.random.rand(self.ndims()) for _ in range(len(texts))] ... >>> registry.get("my-embedding-function") ctSN) __REGISTRY__)clss `/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/lancedb/embeddings/registry.py get_instancez&EmbeddingFunctionRegistry.get_instance,sc\i|_i|_t|_t|_yr) _functions _variablesset_allowed_remote_code_allowed_endpointsselfs r__init__z"EmbeddingFunctionRegistry.__init__0s$=@U!,/ErNaliascfd}|S)a" This creates a decorator that can be used to register an EmbeddingFunction. Parameters ---------- alias : Optional[str] a human friendly name for the embedding function. If not provided, the class name will be used. ct|ts tdxs |j}|jvr!rj||urt |d|j|<||_|S)Nz'Must be a subclass of EmbeddingFunctionz was already registered) issubclassr TypeError__name__rKeyError%__embedding_function_registry_alias__)rkeyrrs r decoratorz5EmbeddingFunctionRegistry.register..decoratorEswc#45 IJJ'3<??#&DOOC 9z5EmbeddingFunctionRegistry._restore..sC c"$*A*AA'se,--0G1 6 s*-) function_namereasonsr))r2_resolve_variablesmodel_constructremote_code_model_remote_code_allowedrboolvalidate_python!_legacy_trust_remote_code_defaultappendextend endpointsr tuple_original_args)rr.rCrresolvedproberMr9allowedexplicit requested model_namerevisionfuncs` rr>z"EmbeddingFunctionRegistry._restoresL hhtn))$/###/h/'')  /d//7G*h6H'-==01  AAC ,5,A'H( )X',$ H,ZLl#88B~F (|1.  (   -"E'ND??D" rr9r`cZ|jj|| tfy|fy)ao Trust the code in a model repository when an embedding function that loads it is restored from table metadata. Table metadata cannot enable `trust_remote_code` by itself, since anyone who can write to the table could then run code in this process. Embedding functions created in this process with `trust_remote_code=True` are allowed automatically. Parameters ---------- model : str The model name, such as `"nomic-ai/nomic-embed-text-v1.5"`. revision : Optional[str] The revision the table's embedding function must pin. Pinning a commit hash is recommended, since the code at a branch or tag can change after you review it. If None, any revision is trusted, including pull request refs that anyone can push to the repository. N)radd _ANY_REVISIONrr9r`s rallow_remote_codez+EmbeddingFunctionRegistry.allow_remote_codes0( !!%% X%5M D ;C D rrJcL|jjt|y)a Trust a custom endpoint, such as an OpenAI `base_url` or Ollama `host`, when an embedding function that uses it is restored from table metadata. Table metadata cannot pick an endpoint by itself, since it could send your data and API keys to a server you don't control. Embedding functions created in this process are allowed to use their endpoints automatically. Parameters ---------- url : str The endpoint URL. Trailing slashes are ignored. N)rrcrH)rrJs rallow_endpointz(EmbeddingFunctionRegistry.allow_endpoints ##N3$78rcN|tf|jvxs||f|jvSr)rdrres rrQz.EmbeddingFunctionRegistry._remote_code_alloweds<}%)B)BB'  G  & &G' 'rrac|j}|(t|ddr|jj||j D]}|j |y)NrEF)rPgetattrrrcrWrh)rrar9rJs r_trustz EmbeddingFunctionRegistry._trusts\&&(  /BE!J  % % ) )% 0>># %C    $ %rc|j}t|tr |j}n!t |d|j j }|j}|||j|jdS)z Convert the given embedding function and source / vector column configs into a config dictionary that can be serialized into arrow metadata r&)r.r9r8r7) r: isinstancer rLrk __class__r$safe_model_dumpr8r7)rconfrar. json_datas rfunction_to_metadataz.EmbeddingFunctionRegistry.function_to_metadatasq }} d6 7%%D=t~~?V?VD((* !//!//   rc|t|dk(ry|Dcgc]}|j|}}tj|dj d}d|iScc}w)z Convert a list of embedding functions and source / vector configs into a config dictionary that can be serialized into arrow metadata Nr)indentr6r5)lenrsr;dumpsencode)r func_listrarrr3s rget_table_metadataz,EmbeddingFunctionRegistry.get_table_metadata sd  I! 3AJKT..t4K K::i299'B%x00 LsAvaluec@d|vr td||j|<y)aS Set a variable. These can be accessed in embedding configuration using the syntax `$var:variable_name`. If they are not set, an error will be thrown letting you know which variable is missing. If you want to supply a default value, you can add an additional part in the configuration like so: `$var:variable_name:default_value`. Default values can be used for runtime configurations that are not sensitive, such as whether to use a GPU for inference. The name must not contain a colon. Default values can contain colons. :z$Variable names cannot contain colonsN) ValueErrorr)rr.r|s rset_varz!EmbeddingFunctionRegistry.set_vars% $;CD D %rc |j|S)z! Get a variable. )rr1s rget_varz!EmbeddingFunctionRegistry.get_var*st$$rr)rqr )r$ __module__ __qualname____doc__ classmethodrrrstrr*r-rr r2rbytesrBdictr>rfrhrRrQrlrsr{rrr)rrrrs362hsm< % %%6 7 %%  eUl!34% c,, -% N1S111B1f s hsm t 09#9$9 '#'#'4' %,%% ( 1&C&&& %C%C%rrrJr/c$|jdS)N/)rstrip)rJs rrHrH4s ::c?rcHtjj|Sr)rrr*)r.s rr*r*=s  $ $ & / / 55rc*tjS)a8 Utility function to get the global instance of the registry Returns ------- EmbeddingFunctionRegistry The global registry instance Examples -------- from lancedb.embeddings import get_registry registry = get_registry() openai = registry.get("openai").create() )rrr)rr get_registryrAs  $ $ &&r)r;typingrrrrrpydanticrbaser r r robjectrdrrHrr*rr)rrrsb 33 ^%^%B   )* 6'/'r