WIP litellm integration

projects/extension/ai/litellm.py

@@ -0,0 +1,35 @@
+import litellm
+from typing import Optional, Generator
+
+
+def embed(
+    model: str,
+    input: list[str],
+    api_key: str,
+    user: Optional[str] = None,
+    dimensions: Optional[int] = None,
+    timeout: Optional[int] = None,
+    api_base: Optional[str] = None,
+    api_version: Optional[str] = None,
+    api_type: Optional[str] = None,
+    organization: Optional[str] = None,
+    **kwargs,
+) -> Generator[tuple[int, list[float]], None, None]:
+    if organization is not None:
+        litellm.organization = organization
+    response = litellm.embedding(
+        model=model,
+        input=input,
+        user=user,
+        dimensions=dimensions,
+        timeout=timeout,
+        api_type=api_type,
+        api_key=api_key,
+        api_base=api_base,
+        api_version=api_version,
+        **kwargs,
+    )
+    if not hasattr(response, "data"):
+        return None
+    for idx, obj in enumerate(response["data"]):
+        yield idx, obj["embedding"]

projects/extension/ai/secrets.py

@@ -24,27 +24,21 @@ def remove_secret_from_cache(sd_cache: dict[str, str], secret_name: str):
 
 def get_secret(
     plpy,
-    secret: Optional[str],
-    secret_name: Optional[str],
-    secret_name_default: str,
-    sd_cache: Optional[dict[str, str]],
-) -> str:
+    secret: Optional[str] = None,
+    secret_name: Optional[str] = None,
+    secret_name_default: Optional[str] = None,
+    sd_cache: Optional[dict[str, str]] = None,
+) -> str | None:
     if secret is not None:
         return secret
 
     if secret_name is None:
         secret_name = secret_name_default
 
     if secret_name is None or secret_name == "":
-        plpy.error("secret_name is required")
-
-    secret = reveal_secret(plpy, secret_name, sd_cache)
-    if secret is None:
-        plpy.error(f"missing {secret_name} secret")
-        # This line should never be reached, but it's here to make the type checker happy.
-        return ""
+        return None
 
-    return secret
+    return reveal_secret(plpy, secret_name, sd_cache)
 
 
 def check_secret_permissions(plpy, secret_name: str) -> bool:

projects/extension/requirements.txt

@@ -1,8 +1,10 @@
-openai==1.44.0
+openai==1.56.0
 tiktoken==0.7.0
 ollama==0.2.1
 anthropic==0.29.0
 cohere==5.5.8
 backoff==2.2.1
 voyageai==0.3.1
 datasets==3.1.0
+litellm==1.55.4
+google-cloud-aiplatform==1.74.0 # required for vertexAI (don't know why litellm doesn't include this)

projects/extension/sql/idempotent/008-embedding.sql

@@ -74,6 +74,30 @@ $func$ language plpgsql immutable security invoker
 set search_path to pg_catalog, pg_temp
 ;
 
+-------------------------------------------------------------------------------
+-- embedding_litellm
+create or replace function ai.embedding_litellm
+( model pg_catalog.text
+, dimensions pg_catalog.int4
+, api_key_name pg_catalog.text default null
+, extra_options pg_catalog.jsonb default null
+) returns pg_catalog.jsonb
+as $func$
+begin
+    return json_object
+    ( 'implementation': 'litellm'
+    , 'config_type': 'embedding'
+    , 'model': model
+    , 'dimensions': dimensions
+    , 'api_key_name': api_key_name
+    , 'extra_options': extra_options
+    absent on null
+    );
+end
+$func$ language plpgsql immutable security invoker
+set search_path to pg_catalog, pg_temp
+;
+
 -------------------------------------------------------------------------------
 -- _validate_embedding
 create or replace function ai._validate_embedding(config pg_catalog.jsonb) returns void
@@ -98,6 +122,8 @@ begin
             -- ok
         when 'voyageai' then
             -- ok
+        when 'litellm' then
+            -- ok
         else
             if _implementation is null then
                 raise exception 'embedding implementation not specified';

projects/extension/sql/idempotent/017-litellm.sql

@@ -0,0 +1,58 @@
+-------------------------------------------------------------------------------
+-- litellm_embed
+-- generate an embedding from a text value
+create or replace function ai.litellm_embed
+( model text
+, input_text text
+, api_key text default null
+, api_key_name text default null
+, extra_options jsonb default null
+) returns @extschema:[email protected]
+as $python$
+    #ADD-PYTHON-LIB-DIR
+    import ai.litellm
+    import ai.secrets
+    options = {}
+    if extra_options is not None:
+        import json
+        options = {k: v for k, v in json.loads(extra_options).items()}
+
+    api_key_resolved = ai.secrets.get_secret(plpy, api_key, api_key_name, None, SD)
+    for tup in ai.litellm.embed(model, [input_text], api_key=api_key_resolved, **options):
+        return tup[1]
+$python$
+language plpython3u immutable parallel safe security invoker
+set search_path to pg_catalog, pg_temp
+;
+
+-------------------------------------------------------------------------------
+-- litellm_embed
+-- generate embeddings from an array of text values
+create or replace function ai.litellm_embed
+( model text
+, input_texts text[]
+, api_key text default null
+, api_key_name text default null
+, extra_options jsonb default null
+) returns table
+( "index" int
+, embedding @extschema:[email protected]
+)
+as $python$
+    #ADD-PYTHON-LIB-DIR
+    import ai.litellm
+    import ai.secrets
+    options = {}
+    if extra_options is not None:
+        import json
+        options = {k: v for k, v in json.loads(extra_options).items()}
+
+    plpy.log("options", options)
+
+    api_key_resolved = ai.secrets.get_secret(plpy, api_key, api_key_name, None, SD)
+    for tup in ai.litellm.embed(model, input_texts, api_key=api_key_resolved, **options):
+        yield tup
+$python$
+language plpython3u immutable parallel safe security invoker
+set search_path to pg_catalog, pg_temp
+;

projects/extension/tests/contents/output16.expected

@@ -34,6 +34,8 @@ CREATE EXTENSION
  function ai.indexing_diskann(integer,text,integer,integer,double precision,integer,integer,boolean)
  function ai.indexing_hnsw(integer,text,integer,integer,boolean)
  function ai.indexing_none()
+ function ai.litellm_embed(text,text,text,text,jsonb)
+ function ai.litellm_embed(text,text[],text,text,jsonb)
  function ai.load_dataset_multi_txn(text,text,text,name,name,text,jsonb,integer,integer,integer,jsonb)
  function ai.load_dataset(text,text,text,name,name,text,jsonb,integer,integer,jsonb)
  function ai.ollama_chat_complete(text,jsonb,text,text,jsonb)
@@ -94,7 +96,7 @@ CREATE EXTENSION
  table ai.vectorizer_errors
  view ai.secret_permissions
  view ai.vectorizer_status
-(90 rows)
+(92 rows)
 
                                     Table "ai._secret_permissions"
  Column | Type | Collation | Nullable | Default | Storage  | Compression | Stats target | Description 

projects/extension/tests/contents/output17.expected

@@ -34,6 +34,8 @@ CREATE EXTENSION
  function ai.indexing_diskann(integer,text,integer,integer,double precision,integer,integer,boolean)
  function ai.indexing_hnsw(integer,text,integer,integer,boolean)
  function ai.indexing_none()
+ function ai.litellm_embed(text,text,text,text,jsonb)
+ function ai.litellm_embed(text,text[],text,text,jsonb)
  function ai.load_dataset_multi_txn(text,text,text,name,name,text,jsonb,integer,integer,integer,jsonb)
  function ai.load_dataset(text,text,text,name,name,text,jsonb,integer,integer,jsonb)
  function ai.ollama_chat_complete(text,jsonb,text,text,jsonb)
@@ -108,7 +110,7 @@ CREATE EXTENSION
  type ai.vectorizer_status[]
  view ai.secret_permissions
  view ai.vectorizer_status
-(104 rows)
+(106 rows)
 
                                     Table "ai._secret_permissions"
  Column | Type | Collation | Nullable | Default | Storage  | Compression | Stats target | Description 

projects/extension/tests/privileges/function.expected

@@ -220,6 +220,14 @@
  f       | bob   | execute   | no      | ai     | indexing_none()
  f       | fred  | execute   | no      | ai     | indexing_none()
  f       | jill  | execute   | YES     | ai     | indexing_none()
+ f       | alice | execute   | YES     | ai     | litellm_embed(model text, input_text text, api_key text, api_key_name text, extra_options jsonb)
+ f       | bob   | execute   | no      | ai     | litellm_embed(model text, input_text text, api_key text, api_key_name text, extra_options jsonb)
+ f       | fred  | execute   | no      | ai     | litellm_embed(model text, input_text text, api_key text, api_key_name text, extra_options jsonb)
+ f       | jill  | execute   | YES     | ai     | litellm_embed(model text, input_text text, api_key text, api_key_name text, extra_options jsonb)
+ f       | alice | execute   | YES     | ai     | litellm_embed(model text, input_texts text[], api_key text, api_key_name text, extra_options jsonb)
+ f       | bob   | execute   | no      | ai     | litellm_embed(model text, input_texts text[], api_key text, api_key_name text, extra_options jsonb)
+ f       | fred  | execute   | no      | ai     | litellm_embed(model text, input_texts text[], api_key text, api_key_name text, extra_options jsonb)
+ f       | jill  | execute   | YES     | ai     | litellm_embed(model text, input_texts text[], api_key text, api_key_name text, extra_options jsonb)
  f       | alice | execute   | YES     | ai     | load_dataset(name text, config_name text, split text, schema_name name, table_name name, if_table_exists text, field_types jsonb, batch_size integer, max_batches integer, kwargs jsonb)
  f       | bob   | execute   | no      | ai     | load_dataset(name text, config_name text, split text, schema_name name, table_name name, if_table_exists text, field_types jsonb, batch_size integer, max_batches integer, kwargs jsonb)
  f       | fred  | execute   | no      | ai     | load_dataset(name text, config_name text, split text, schema_name name, table_name name, if_table_exists text, field_types jsonb, batch_size integer, max_batches integer, kwargs jsonb)
@@ -328,5 +336,5 @@
  f       | bob   | execute   | no      | ai     | voyageai_embed(model text, input_texts text[], input_type text, api_key text, api_key_name text)
  f       | fred  | execute   | no      | ai     | voyageai_embed(model text, input_texts text[], input_type text, api_key text, api_key_name text)
  f       | jill  | execute   | YES     | ai     | voyageai_embed(model text, input_texts text[], input_type text, api_key text, api_key_name text)
-(328 rows)
+(336 rows)