Implement URI lookup for AI embedding models.

[dnwpark]

Related #8740

Allows the definition of ai embedding indexes using a uri:

  type Astronomy {
    content: str;
    deferred index ext::ai::index(embedding_model := 'openai:text-embedding-3-small')
      on (.content);
  }

When parsing a URI, the current schema is checked for a matching model type. If none exists, a reference json is checked and a new type is created if necessary.

For example, if text-embedding-3-small did not exist in the schema, the type __ext_generated_types__::ai_embedding_openai_text-embedding-3-small with the appropriate annotations would be generated. This type is automatically cleaned up when the index is deleted.

[elprans]

Is this just to keep the reference? I wonder if you can formulate the generated ref injection via find_extra_refs callback in compile_expr_field instead?

[dnwpark]

Ah, I am not actually compiling anything into an expression, so I'm not sure that would work.

First, I manually create a qlast.CreateObjectType for the new model type:

gel/edb/schema/indexes.py

Line 1272 in 5f8e00b

model_ast = qlast.CreateObjectType(

I compile it:

gel/edb/schema/indexes.py

Lines 1335 to 1339 in 5f8e00b

	model_cmd = sd.compile_ddl(
	schema,
	model_ast,
	)
	model_cmd.set_attribute_value('is_schema_generated', True)

Then, return to IndexCommand._handle_ai_index_op and make it a prerequisite of the CreateIndex/AlterIndex command:

gel/edb/schema/indexes.py

Line 1025 in 5f8e00b

self.add_prerequisite(model_cmd)

However, to get the typeshell of the type produced by the command, I compile the command a second time and actually apply it to the schema:

gel/edb/schema/indexes.py

Lines 1341 to 1355 in 5f8e00b

	# Doing this since using model_cmd or a copy doesn't work
	dummy_model_cmd = sd.compile_ddl(
	schema,
	model_ast,
	)
	dummy_model_cmd.set_attribute_value('is_schema_generated', True)
	new_schema = dummy_model_cmd.apply(schema, context)
	model_type = new_schema.get(model_typename, None, type=s_types.Type)
	assert model_type is not None
	model_typeshell = so.ObjectShell(
	name=model_typename,
	schemaclass=type(model_type),
	)
	return model_cmd, model_typeshell

[aljazerzen]

We have similar fields from_alias and from_global. Maybe they can be reused?

[dnwpark]

from_alias has a lot of associated logic with links and views. I thought it best to avoid adding more meanings to it. @msullivan will hopefully have more insight

[aljazerzen]

I understand about half of this, but also don't really understand what it does - even after reading the tests. From what I gather, it allows using AI models from the JSON file in the user schema?

I can spend more time here, but I'd need an explanation.

[dnwpark]

@aljazerzen The issue should explain it #8740. If there is anything in particular you don't understand please ask! I will also add some more comments.

[anbuzin]

URI selection seems to work great interface-wise. The LSP is not happy though:

undefined embedding model: no subtype of ext::ai::EmbeddingModel is annotated as 'openai:text-embedding-3-small'

There were some issues with the Ollama implementation, but that's not the PR to talk about them, so I'll do more research and file them separately.

[dnwpark]

@anbuzin What sort of issues were you having with ollama? Right now, not all models are in the reference json yet so that might be it.

[anbuzin]

@dnwpark It was more like, you have to insert a config with an empty string for a secret, and I couldn't get the RAG to work with llama3.2, but I haven't had a chance to figure out exactly why

[dnwpark]

@anbuzin Ah yeah, you need to config the provider no matter what. In theory, you could override api_url to have ollama to run on a server.

[msullivan]

We need tests where we add new entries to the manifest, after having done the initial creation.

We should probably only check the file once per server startup (maybe? or in the eventual case maybe once per 24 hours of something).

Since this version doesn't support the http loading anyway, I think the testing approach should be to support passing the file name via an environment variable. Then we can create a temp test dir, start up the server in that directory with start_edgedb_server, do some work, stop it, populate a new updated file, start the server pointing at that instead, and then do the rest of the testing.

Does that make sense?

[msullivan]

I've concluded that I don't like this approach. I think that this will be cleaner if we take a "JSON-forward" approach, where instead of driving the extension worker based on querying the introspection schema, we drive it based on an in-memory datastructure driven by the JSON. It should be pretty easy to produce a data structure for the db-configured stuff with an introspection query; probably we could directly produce the JSON directly using json_object_pack, if we wanted...

One thing that we will need to do either in that new approach or if we wanted to continue this approach, is to cache the last downloaded version of the file. Otherwise, if the fetch fails on server startup or whatever, there will be a degradation of service. In the JSON-forward approach it would be worse, but even in this schema-based approach it would prevent running migration create.

[msullivan]

I think we should pass this in as a mapping to the data. Or if we want to generalize things further, to some abstract interface type for retrieving data?

[msullivan]

Hm... let me dig into something a bit more. I had forgotten that we actually copy the annotations onto the indexes themselves... but I can't remember why?

Edit: Ah, because it's used to make the data more easily accessible in delta_ext_ai.py and in the introspection queries?

gel/edb/schema/indexes.py

Lines 1468 to 1470 in fe28a9f

	# Copy ext::ai:: annotations declared on the model specified
	# by the `embedding_model` kwarg. This is necessary to avoid
	# expensive lookups later where the index is used.

I don't really believe that the later lookups would be particularly expensive without it, though.

[msullivan]

Here's a query that will generate the exact same JSON I think, from introspecting:

WITH ai_objs := (
  SELECT schema::ObjectType FILTER 'ext::ai::EmbeddingModel' IN .ancestors.name
)
SELECT json_object_pack(
  FOR obj IN ai_objs
  FOR j IN json_object_pack(
    FOR a IN obj.annotations
    SELECT (
        re_replace('^embedding_model_', '',
          re_replace('^ext::ai::', '', a.name),
        ),
        if a.name like 'ext::ai::model_%'
        then <json>a@value else to_json(a@value)
    )
    FILTER a.name LIKE 'ext::ai::%'
  )
  SELECT (<str>j['model_name'], j)
);

Though maybe you'd be better off just doing

SELECT schema::ObjectType {
    annotations: {name, @value},
}
FILTER 'ext::ai::EmbeddingModel' IN .ancestors.name;

and postprocessing in python :)