embedding queries (#816)

* Adds embedding based capabilites * dramatically updates chat API capabilities --------- Co-authored-by: Edward Louth <edward@louths.com>
grai-io · Dec 15, 2023 · 248e66b · 248e66b
1 parent 2d3ab13
commit 248e66b
Show file tree

Hide file tree

Showing 50 changed files with 1,648 additions and 1,104 deletions.
diff --git a/.github/workflows/cli-ci.yml b/.github/workflows/cli-ci.yml
@@ -96,7 +96,7 @@ jobs:
 
     services:
       postgres:
-        image: postgres:latest
+        image: ankane/pgvector:latest
         env:
           POSTGRES_USER: grai
           POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/client-ci.yml b/.github/workflows/client-ci.yml
@@ -91,7 +91,7 @@ jobs:
         working-directory: ${{ env.project_dir }}
     services:
       postgres:
-        image: postgres:latest
+        image: ankane/pgvector:latest
         env:
           POSTGRES_USER: grai
           POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/graph-ci.yml b/.github/workflows/graph-ci.yml
@@ -44,7 +44,7 @@ jobs:
         working-directory: ${{ env.project_dir }}
     services:
       postgres:
-        image: postgres:latest
+        image: ankane/pgvector:latest
         env:
           POSTGRES_USER: grai
           POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/integration-bigquery-ci.yml b/.github/workflows/integration-bigquery-ci.yml
@@ -46,7 +46,7 @@ jobs:
 #        working-directory: ${{ env.project_dir }}
 #    services:
 #      postgres:
-#        image: postgres:latest
+#        image: ankane/pgvector:latest
 #        env:
 #          POSTGRES_USER: grai
 #          POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/integration-fivetran-ci.yml b/.github/workflows/integration-fivetran-ci.yml
@@ -39,7 +39,7 @@ jobs:
         working-directory: ${{ env.project_dir }}
     services:
       postgres:
-        image: postgres:latest
+        image: ankane/pgvector:latest
         env:
           POSTGRES_USER: grai
           POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/integration-flat_file-ci.yml b/.github/workflows/integration-flat_file-ci.yml
@@ -46,7 +46,7 @@ jobs:
         working-directory: ${{ env.project_dir }}
     services:
       postgres:
-        image: postgres:latest
+        image: ankane/pgvector:latest
         env:
           POSTGRES_USER: grai
           POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/integration-postgres-ci.yml b/.github/workflows/integration-postgres-ci.yml
@@ -48,7 +48,7 @@ jobs:
         working-directory: ${{ env.project_dir }}
     services:
       postgres:
-        image: postgres:latest
+        image: ankane/pgvector:latest
         env:
           POSTGRES_USER: grai
           POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/integration-snowflake-ci.yml b/.github/workflows/integration-snowflake-ci.yml
@@ -48,7 +48,7 @@ jobs:
 #        working-directory: ${{ env.project_dir }}
 #    services:
 #      postgres:
-#        image: postgres:latest
+#        image: postgankane/pgvectorres:latest
 #        env:
 #          POSTGRES_USER: grai
 #          POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -77,7 +77,7 @@ jobs:
 
     services:
       postgres:
-        image: postgres:latest
+        image: ankane/pgvector:latest
         env:
           POSTGRES_USER: grai
           POSTGRES_PASSWORD: grai

diff --git a/.github/workflows/server-ci.yml b/.github/workflows/server-ci.yml
@@ -52,7 +52,7 @@ jobs:
 
     services:
       postgres:
-        image: postgres:latest
+        image: ankane/pgvector:latest
         env:
           POSTGRES_USER: ${{ env.DB_USER }}
           POSTGRES_PASSWORD: ${{ env.DB_PASSWORD }}

diff --git a/docs/pages/deployment/docker.mdx b/docs/pages/deployment/docker.mdx
@@ -84,7 +84,7 @@ version: "3.7"
 
 services:
   db:
-    image: postgres:14.3-alpine
+    image: ankane/pgvector
     environment:
       - POSTGRES_USER=grai
       - POSTGRES_PASSWORD=grai

diff --git a/docs/pages/deployment/kubernetes/kustomize.mdx b/docs/pages/deployment/kubernetes/kustomize.mdx
@@ -50,7 +50,7 @@ spec:
             - containerPort: 8000
               protocol: TCP
         - name: db-pod
-          image: postgres:14.3-alpine
+          image: ankane/pgvector
           env:
             - name: POSTGRES_USER
               value: grai

diff --git a/examples/deployment/docker-compose/complete/docker-compose.yml b/examples/deployment/docker-compose/complete/docker-compose.yml
@@ -2,7 +2,7 @@ version: "3.7"
 
 services:
   db:
-    image: postgres:14.3-alpine
+    image: ankane/pgvector
     extra_hosts:
       - "host.docker.internal:host-gateway"
     environment:

diff --git a/examples/deployment/docker-compose/demo/docker-compose.yaml b/examples/deployment/docker-compose/demo/docker-compose.yaml
@@ -2,7 +2,7 @@ version: "3.7"
 
 services:
   db:
-    image: postgres:14.3-alpine
+    image: ankane/pgvector
     environment:
       - POSTGRES_USER=grai
       - POSTGRES_PASSWORD=grai

diff --git a/examples/deployment/docker-compose/minimal/docker-compose.yml b/examples/deployment/docker-compose/minimal/docker-compose.yml
@@ -2,7 +2,7 @@ version: "3.7"
 
 services:
   db:
-    image: postgres:14.3-alpine
+    image: ankane/pgvector
     environment:
       - POSTGRES_USER=grai
       - POSTGRES_PASSWORD=grai

diff --git a/examples/deployment/k8s/deployment.yaml b/examples/deployment/k8s/deployment.yaml
@@ -34,7 +34,7 @@ spec:
           - containerPort: 8000
             protocol: TCP
         - name: db-pod
-          image: postgres:14.3-alpine
+          image: ankane/pgvector
           env:
           - name: POSTGRES_USER
             value: grai

diff --git a/grai-client/docker-compose.yml b/grai-client/docker-compose.yml
@@ -32,7 +32,7 @@ services:
       start_period: 20s
 
   db:
-    image: postgres:14.3-alpine
+    image: ankane/pgvector
     environment:
       - POSTGRES_USER=grai
       - POSTGRES_PASSWORD=grai

diff --git a/grai-integrations/source-dbt-cloud/docker-compose.yml b/grai-integrations/source-dbt-cloud/docker-compose.yml
@@ -3,7 +3,7 @@ version: "3.7"
 
 services:
   postgres:
-    image: postgres:latest
+    image: ankane/pgvector
     environment:
       - POSTGRES_USER=grai
       - POSTGRES_PASSWORD=grai

diff --git a/grai-integrations/source-dbt/docker-compose.yml b/grai-integrations/source-dbt/docker-compose.yml
@@ -3,7 +3,7 @@ version: "3.7"
 
 services:
   postgres:
-    image: postgres:latest
+    image: ankane/pgvector
     environment:
       - POSTGRES_USER=grai
       - POSTGRES_PASSWORD=grai

diff --git a/grai-integrations/source-metabase/docker-compose.yml b/grai-integrations/source-metabase/docker-compose.yml
@@ -2,7 +2,7 @@ version: '3.8'
 
 services:
   db:
-    image: postgres:13
+    image: ankane/pgvector
     environment:
       POSTGRES_PASSWORD: mysecretpassword
       POSTGRES_USER: metabase

diff --git a/grai-integrations/source-mysql/docker-compose.yml b/grai-integrations/source-mysql/docker-compose.yml
@@ -22,7 +22,7 @@ services:
 #      retries: 10
 #      start_period: 20s
 #  db:
-#    image: postgres:latest
+#    image: ankane/pgvector
 #    environment:
 #      - POSTGRES_USER=grai
 #      - POSTGRES_PASSWORD=grai

diff --git a/grai-integrations/source-postgres/docker-compose.yml b/grai-integrations/source-postgres/docker-compose.yml
@@ -2,7 +2,7 @@ version: "3.7"
 
 services:
   postgres:
-    image: "postgres:latest"
+    image: ankane/pgvector
     volumes:
      - ./dev-database/schemas:/docker-entrypoint-initdb.d
     environment:

diff --git a/grai-server/app/ai_build_script.py b/grai-server/app/ai_build_script.py
@@ -2,7 +2,7 @@
 
 import tiktoken
 
-supported_models = ["gpt-4", "gpt-4-32k", "gpt-3.5-turbo"]
+supported_models = ["gpt-4", "gpt-4-32k", "gpt-3.5-turbo", "gpt-4-1106-preview"]
 
 cache_dir = os.environ.get("TIKTOKEN_CACHE_DIR", None)
 print("CACHE DIR: ", cache_dir)

diff --git a/grai-server/app/connections/adapters/schemas.py b/grai-server/app/connections/adapters/schemas.py
@@ -1,5 +1,5 @@
 import pprint
-from typing import Any, List, Literal, Optional, Sequence, Type
+from typing import Any, List, Literal, Optional, Sequence, Type, TypeVar
 from uuid import UUID
 
 from django.db.models import Q
@@ -20,6 +20,9 @@
 from lineage.models import Edge, Node, Source
 from workspaces.models import Organisation, Workspace
 
+T = TypeVar("T")
+R = TypeVar("R")
+
 
 def get_data_source_models(data_sources: Sequence[SourceSpec], workspace: Workspace) -> List[Source]:
     uuids = [source_id for source_id in data_sources if isinstance(source_id, UUID)]
@@ -198,7 +201,10 @@ def source_model_to_source_schema(model: Source, schema_type: Literal["SourceV1"
 @model_to_schema.register
 def node_model_to_node_v1_schema(model: Node, schema_type: Literal["NodeV1"]) -> NodeV1:
     # TODO: Add data_sources
-    return NodeV1.from_spec({**model.__dict__, "data_sources": []})
+
+    data_sources: list[SourceV1] = model_to_schema(model.data_sources.all(), "SourceV1")
+    result = NodeV1.from_spec({**model.__dict__, "data_sources": [source.spec for source in data_sources]})
+    return result
 
 
 @model_to_schema.register
@@ -220,11 +226,11 @@ def edge_model_to_edge_v1_schema(model: Edge, schema_type: Literal["EdgeV1"]) ->
 
 
 @model_to_schema.register
-def sequence_model_to_sequence_v1_schema(models: list | tuple, schema_type: str) -> list | tuple:
-    iter = (model_to_schema(model, schema_type) for model in models)
-    return type(models)(iter)
+def sequence_model_to_sequence_v1_schema(models: list[T] | tuple[T], schema_type: str) -> list[R]:
+    result = list(model_to_schema(model, schema_type) for model in models)
+    return result
 
 
 @model_to_schema.register
-def queryset_to_sequence_v1_schema(models: QuerySet, schema_type: str) -> list:
+def queryset_to_sequence_v1_schema(models: QuerySet[T], schema_type: str) -> list[R]:
     return [model_to_schema(model, schema_type) for model in models]
diff --git a/grai-server/app/connections/tests/test_tasks.py b/grai-server/app/connections/tests/test_tasks.py
@@ -265,12 +265,7 @@ def test_run_update_server_postgres_wrong_password(self, test_workspace, test_po
 
         assert run.status == "error"
         assert run.metadata["error"] == "Incorrect password"
-        assert (
-            run.metadata["message"]
-            == 'connection to server at "localhost" (127.0.0.1), port 5432 failed: FATAL:  password authentication failed for user "grai"\n'
-            or run.metadata["message"]
-            == 'connection to server at "127.0.0.1", port 5432 failed: FATAL:  password authentication failed for user "grai"\n'
-        )
+        assert run.metadata["message"].endswith('FATAL:  password authentication failed for user "grai"\n')
 
     def test_run_update_server_postgres_no_database(self, test_workspace, test_postgres_connector, test_source):
         connection = Connection.objects.create(
@@ -294,12 +289,7 @@ def test_run_update_server_postgres_no_database(self, test_workspace, test_postg
 
         assert run.status == "error"
         assert run.metadata["error"] == "Missing permission"
-        assert (
-            run.metadata["message"]
-            == 'connection to server at "localhost" (127.0.0.1), port 5432 failed: FATAL:  database "wrong" does not exist\n'
-            or run.metadata["message"]
-            == 'connection to server at "127.0.0.1", port 5432 failed: FATAL:  database "wrong" does not exist\n'
-        )
+        assert run.metadata["message"].endswith('FATAL:  database "wrong" does not exist\n')
 
     def test_run_update_server_no_connector(self, test_workspace, test_connector, test_source):
         connection = Connection.objects.create(

diff --git a/grai-server/app/grAI/authentication.py b/grai-server/app/grAI/authentication.py
@@ -31,7 +31,6 @@ async def __call__(self, scope, receive, send):
             raise PermissionDeniedError("You do not have permission to access this resource.")
 
         scope.setdefault("metadata", {})
-        scope["metadata"]["workspace_id"] = workspace_id
         scope["metadata"]["membership"] = membership
         return await super().__call__(scope, receive, send)