Skip to content

Support pgvector extension in codegen #446

@comfuture

Description

@comfuture

Describe the bug

edgedb python doesn't seems support pgvector yet.
When I try to create codegen for pgvector, it throws an error.

Reproduction

Simple edgeql that uses pgvector:

with
  vec as module ext::pgvector,
  target := <vec1536>$vector,
  threshold := <float64>$threshold

select Document {
  *,
  dist := vec::cosine_distance(target, .embedding)
}
filter .dist < threshold
order by .dist empty last
limit 5

Try to create code by edgedb-py command:

edgedb-py

Error message

$ edgedb-py
Found EdgeDB project: /workspace
Processing /workspace/queries/search_document.edgeql
Generating /workspace/queries/search_document_async_edgeql.py
Traceback (most recent call last):
  File "/home/vscode/.local/bin/edgedb-py", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/vscode/.local/lib/python3.11/site-packages/edgedb/codegen/cli.py", line 99, in main
    generator.Generator(args).run()
  File "/home/vscode/.local/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 198, in run
    self._generate_files(suffix)
  File "/home/vscode/.local/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 232, in _generate_files
    content = self._generate(name, query, dr)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/vscode/.local/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 319, in _generate
    out_type = self._generate_code(dr.output_type, name_hint)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/vscode/.local/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 452, in _generate_code
    el_code = self._generate_code_with_cardinality(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/vscode/.local/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 524, in _generate_code_with_cardinality
    rv = self._generate_code(type_, name_hint)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/vscode/.local/lib/python3.11/site-packages/edgedb/codegen/generator.py", line 434, in _generate_code
    value = TYPE_MAPPING[base_type_name]
            ~~~~~~~~~~~~^^^^^^^^^^^^^^^^
KeyError: 'ext::pgvector::vector'

Expected behavior
Should be able to generate python file named search_document_async_edgeql.py

Versions (please complete the following information):

  • OS: Linux 5ca54fafbd62 5.15.0-1040-azure #47-Ubuntu SMP Thu Jun 1 19:38:24 UTC 2023 x86_64 GNU/Linux
  • EdgeDB version: 3.0+304325b
  • EdgeDB CLI version: 3.3.0+071876c
  • edgedb-python version: 1.5.0
  • Python version: Python 3.11.4

Additional context
Add any other context about the problem here.

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions