Update ai ext

docs/ai/javascript.rst

@@ -171,3 +171,43 @@ Public methods
         Optional. Variable settings required for the context query.
     :param number context.max_object_count:
         Optional. A maximum number of objects to return from the context query.
+
+.. js:method:: async streamRag( \
+                   message: string, \
+                   context: QueryContext = this.context \
+                 ): AsyncIterable<StreamingMessage> & PromiseLike<Response>
+
+    Can be used in two ways:
+
+    - as **an async iterator** - if you want to process streaming data in 
+        real-time as it arrives, ideal for handling long-running streams.
+
+    - as **a Promise that resolves to a full Response object** - you have 
+        complete control over how you want to handle the stream, this might be 
+        useful when you want to manipulate the raw stream or parse it in a custom way.
+
+    :param string message:
+        Required. The message to be sent to the text generation provider's API.
+    :param string context.query:
+        Required. Specifies an expression to determine the relevant objects and
+        index to serve as context for text generation. You may set this to any
+        expression that produces a set of objects, even if it is not a
+        standalone query.
+    :param string context.variables:
+        Optional. Variable settings required for the context query.
+    :param string context.globals:
+        Optional. Variable settings required for the context query.
+    :param number context.max_object_count:
+        Optional. A maximum number of objects to return from the context query.
+
+.. js:method:: async generateEmbeddings( \
+                   inputs: string[], \
+                   model: string \
+                 ): Promise<number[]>
+
+    Generates embeddings for the array of strings.
+
+    :param string[] inputs:
+        Required. Strings array to generate embeddings for.
+    :param string model:
+        Required. Specifies the AI model to use.

docs/reference/edgeql/tx_start.rst

@@ -61,14 +61,13 @@ Parameters
 The :eql:synopsis:`<transaction-mode>` can be one of the following:
 
 :eql:synopsis:`isolation serializable`
-    All statements of the current transaction can only see data
-    changes committed before the first query or data-modification
-    statement was executed in this transaction.  If a pattern
-    of reads and writes among concurrent serializable
-    transactions would create a situation which could not have
-    occurred for any serial (one-at-a-time) execution of those
-    transactions, one of them will be rolled back with a
-    serialization_failure error.
+    All statements in the current transaction can only see data 
+    changes that were committed before the first query or data 
+    modification statement was executed within this transaction. 
+    If a pattern of reads and writes among concurrent serializable
+    transactions creates a situation that could not have occurred 
+    in any serial (one-at-a-time) execution of those transactions, 
+    one of them will be rolled back with a serialization_failure error.
 
 :eql:synopsis:`read write`
     Sets the transaction access mode to read/write.

edb/lib/ext/ai.edgeql

@@ -140,9 +140,10 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
     create type ext::ai::Config extending cfg::ExtensionConfig {
         create required property indexer_naptime: std::duration {
             set default := <std::duration>'10s';
-            create annotation std::description :=
-                "Specifies the minimum delay between deferred ext::ai::index "
-                ++ "indexer runs on any given branch.";
+            create annotation std::description := '
+                Specifies the minimum delay between runs of the 
+                deferred ext::ai::index indexer on any given branch.
+            ';
         };
 
         create multi link providers: ext::ai::ProviderConfig {
@@ -196,8 +197,7 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
             ext::ai::text_gen_model_context_window := "<must override>";
     };
 
-    # OpenAI models.
-
+    # OpenAI models
     create abstract type ext::ai::OpenAITextEmbedding3SmallModel
         extending ext::ai::EmbeddingModel
     {
@@ -261,7 +261,7 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
             ext::ai::text_gen_model_context_window := "16385";
     };
 
-    create abstract type ext::ai::OpenAIGPT_4_TurboModel
+    create abstract type ext::ai::OpenAIGPT_4_TurboPreviewModel
         extending ext::ai::TextGenerationModel
     {
         alter annotation
@@ -272,7 +272,51 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
             ext::ai::text_gen_model_context_window := "128000";
     };
 
-    # Mistral models.
+    create abstract type ext::ai::OpenAIGPT_4_TurboModel
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "gpt-4-turbo";
+        alter annotation
+            ext::ai::model_provider := "builtin::openai";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+    create abstract type ext::ai::OpenAIGPT_4o_Model
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "gpt-4o";
+        alter annotation
+            ext::ai::model_provider := "builtin::openai";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+    create abstract type ext::ai::OpenAIGPT_4o_MiniModel
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "gpt-4o-mini";
+        alter annotation
+            ext::ai::model_provider := "builtin::openai";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+    create abstract type ext::ai::OpenAIGPT_4_Model
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "gpt-4";
+        alter annotation
+            ext::ai::model_provider := "builtin::openai";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+    # Mistral models
     create abstract type ext::ai::MistralEmbedModel
         extending ext::ai::EmbeddingModel
     {
@@ -296,10 +340,11 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
         alter annotation
             ext::ai::model_provider := "builtin::mistral";
         alter annotation
-            ext::ai::text_gen_model_context_window := "8192";
+            ext::ai::text_gen_model_context_window := "32000";
     };
 
-    create abstract type ext::ai::MistralMediumModel
+    # going to be deprecated shortly
+    create abstract type ext::ai::MistralMediumModel 
         extending ext::ai::TextGenerationModel
     {
         alter annotation
@@ -318,10 +363,90 @@ CREATE EXTENSION PACKAGE ai VERSION '1.0' {
         alter annotation
             ext::ai::model_provider := "builtin::mistral";
         alter annotation
-            ext::ai::text_gen_model_context_window := "8192";
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+    create abstract type ext::ai::Ministral_8B
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "ministral-8b-latest";
+        alter annotation
+            ext::ai::model_provider := "builtin::mistral";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+    create abstract type ext::ai::Ministral_3B
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "ministral-3b-latest";
+        alter annotation
+            ext::ai::model_provider := "builtin::mistral";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+    create abstract type ext::ai::Codestral
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "codestral-latest";
+        alter annotation
+            ext::ai::model_provider := "builtin::mistral";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+
+    # Mistral free models
+    create abstract type ext::ai::Pixtral
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "pixtral-12b-2409";
+        alter annotation
+            ext::ai::model_provider := "builtin::mistral";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+
+    create abstract type ext::ai::MistralNemo
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "open-mistral-nemo";
+        alter annotation
+            ext::ai::model_provider := "builtin::mistral";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "128000";
+    };
+
+    create abstract type ext::ai::CodestralMamba
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "open-codestral-mamba";
+        alter annotation
+            ext::ai::model_provider := "builtin::mistral";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "256000";
+    };
+
+    # Anthropic models
+    create abstract type ext::ai::AnthropicClaude_3_5_SonnetModel
+        extending ext::ai::TextGenerationModel
+    {
+        alter annotation
+            ext::ai::model_name := "claude-3-5-sonnet-20240620";
+        alter annotation
+            ext::ai::model_provider := "builtin::anthropic";
+        alter annotation
+            ext::ai::text_gen_model_context_window := "200000";
     };
 
-    # Anthropic models.
     create abstract type ext::ai::AnthropicClaude3HaikuModel
         extending ext::ai::TextGenerationModel
     {

edb/server/http.py

@@ -488,6 +488,9 @@ class SSEEvent:
         data: str
         id: Optional[str] = None
 
+        def json(self):
+            return json_lib.loads(self.data)
+
     def close(self):
         self._cancel()
 

edb/server/http/src/python.rs

@@ -178,9 +178,13 @@ async fn request_sse(
     let guard = guard((), |_| trace!("Exiting SSE due to cancellation"));
     let response = request(client, url, method, body, headers).await?;
 
-    if response.headers().get("content-type")
-        != Some(&HeaderValue::from_static("text/event-stream"))
-    {
+    let content_type = response.headers().get("content-type");
+    let is_event_stream = content_type
+        .and_then(|v| v.to_str().ok())
+        .map(|s| s.starts_with("text/event-stream"))
+        .unwrap_or(false);
+
+    if !is_event_stream {
         let headers = process_headers(response.headers());
         let status = response.status();
         let body = match response.bytes().await {