Add python binding for vectorstore

gel/_testbase.py

@@ -178,12 +178,12 @@ def _start_cluster(*, cleanup_atexit=True):
                         if line.startswith(b'READY='):
                             break
                     else:
-                        raise RuntimeError('not ready')
+                        raise RuntimeError("not ready")
                 break
             except Exception:
                 time.sleep(1)
         else:
-            raise RuntimeError('server status file not found')
+            raise RuntimeError("server status file not found")
 
         data = json.loads(line.split(b'READY=')[1])
         con_args = dict(host='localhost', port=data['port'])
@@ -203,8 +203,7 @@ def _start_cluster(*, cleanup_atexit=True):
                     ]
                 )
             else:
-                con_args['tls_ca_file'] = data['tls_cert_file']
-
+                con_args["tls_ca_file"] = data["tls_cert_file"]
         client = gel.create_client(password='test', **con_args)
         client.ensure_connected()
         client.execute("""

gel/ai/__init__.py

@@ -1,7 +1,7 @@
 #
-# This source file is part of the EdgeDB open source project.
+# This source file is part of the Gel open source project.
 #
-# Copyright 2024-present MagicStack Inc. and the EdgeDB authors.
+# Copyright 2024-present MagicStack Inc. and the Gel authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,6 +19,21 @@
 from .types import RAGOptions, ChatParticipantRole, Prompt, QueryContext
 from .core import create_rag_client, RAGClient
 from .core import create_async_rag_client, AsyncRAGClient
+from .vectorstore import (
+    create_vstore,
+    create_async_vstore,
+    Record,
+    Vector,
+    SearchResult,
+    EmbeddingModel,
+    AsyncEmbeddingModel,
+)
+from .metadata_filter import (
+    MetadataFilter,
+    CompositeFilter,
+    FilterOperator,
+    FilterCondition,
+)
 
 __all__ = [
     "RAGOptions",
@@ -29,4 +44,15 @@
     "RAGClient",
     "create_async_rag_client",
     "AsyncRAGClient",
+    "create_vstore",
+    "create_async_vstore",
+    "Record",
+    "Vector",
+    "SearchResult",
+    "EmbeddingModel",
+    "AsyncEmbeddingModel",
+    "MetadataFilter",
+    "CompositeFilter",
+    "FilterOperator",
+    "FilterCondition",
 ]

gel/ai/metadata_filter.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Union, List
+from enum import Enum
+
+
+class FilterOperator(str, Enum):
+    EQ = "="
+    NE = "!="
+    GT = ">"
+    LT = "<"
+    GTE = ">="
+    LTE = "<="
+    IN = "in"
+    NOT_IN = "not in"
+    LIKE = "like"
+    ILIKE = "ilike"
+    ANY = "any"
+    ALL = "all"
+    CONTAINS = "contains"
+    EXISTS = "exists"
+
+
+class FilterCondition(str, Enum):
+    AND = "and"
+    OR = "or"
+
+
+@dataclass
+class MetadataFilter:
+    """Represents a single metadata filter condition."""
+
+    key: str
+    value: Union[int, float, str]
+    operator: FilterOperator = FilterOperator.EQ
+
+    def __repr__(self):
+        value = f'"{self.value}"' if isinstance(self.value, str) else self.value
+        return (
+            f'MetadataFilter(key="{self.key}", '
+            f"value={value}, "
+            f'operator="{self.operator.value}")'
+        )
+
+
+@dataclass
+class CompositeFilter:
+    """
+    Allows grouping multiple MetadataFilter instances using AND/OR.
+    """
+
+    filters: List[Union[CompositeFilter, MetadataFilter]]
+    condition: FilterCondition = FilterCondition.AND
+
+    def __repr__(self):
+        return (
+            f'CompositeFilter(condition="{self.condition.value}", '
+            f"filters={self.filters})"
+        )
+
+
+def get_filter_clause(filters: CompositeFilter) -> str:
+    """
+    Get the filter clause for a given CompositeFilter.
+    """
+
+    subclauses = []
+    for filter in filters.filters:
+        subclause = ""
+
+        if isinstance(filter, CompositeFilter):
+            subclause = get_filter_clause(filter)
+        elif isinstance(filter, MetadataFilter):
+            formatted_value = (
+                f'"{filter.value}"'
+                if isinstance(filter.value, str)
+                else filter.value
+            )
+
+            # Simple comparison operators
+            if filter.operator in {
+                FilterOperator.EQ,
+                FilterOperator.NE,
+                FilterOperator.GT,
+                FilterOperator.GTE,
+                FilterOperator.LT,
+                FilterOperator.LTE,
+                FilterOperator.LIKE,
+                FilterOperator.ILIKE,
+            }:
+                subclause = (
+                    f'<typeof {formatted_value}>json_get(.metadata, "{filter.key}") '
+                    f"{filter.operator.value} {formatted_value}"
+                )
+            # casting should be fixed
+            elif filter.operator in {FilterOperator.IN, FilterOperator.NOT_IN}:
+                subclause = (
+                    f"{formatted_value} "
+                    f"{filter.operator.value} "
+                    f'<str>json_get(.metadata, "{filter.key}")'
+                )
+            # casting should be fixed
+            # works only for equality, should be updated to support the rest,
+            # example: select all({1, 2, 3, 4} < 4);
+            elif filter.operator in {FilterOperator.ANY, FilterOperator.ALL}:
+                subclause = (
+                    f"{filter.operator.value} ("
+                    f'<str>json_get(.metadata, "{filter.key}")'
+                    f" = {formatted_value})"
+                )
+
+            elif filter.operator == FilterOperator.EXISTS:
+                subclause = f'exists <str>json_get(.metadata, "{filter.key}")'
+
+            # casting should be fixed
+            # edgeql contains supports different types like range etc which
+            # we don't support here
+            elif filter.operator == FilterOperator.CONTAINS:
+                subclause = (
+                    f'contains (<str>json_get(.metadata, "{filter.key}"), '
+                    f"{formatted_value})"
+                )
+            else:
+                raise ValueError(f"Unknown operator: {filter.operator}")
+
+        subclauses.append(subclause)
+
+    if filters.condition in {FilterCondition.AND, FilterCondition.OR}:
+        filter_clause = f" {filters.condition.value} ".join(subclauses)
+        return (
+            "(" + filter_clause + ")" if len(subclauses) > 1 else filter_clause
+        )
+    else:
+        raise ValueError(f"Unknown condition: {filters.condition}")