diff --git a/agents-api/agents_api/autogen/Docs.py b/agents-api/agents_api/autogen/Docs.py index a7023ddfc..0c1524e48 100644 --- a/agents-api/agents_api/autogen/Docs.py +++ b/agents-api/agents_api/autogen/Docs.py @@ -57,6 +57,13 @@ class Doc(BaseModel): """ Contents of the document """ + embeddings: Annotated[ + list[float] | list[list[float]] | None, + Field(None, json_schema_extra={"readOnly": True}), + ] + """ + Embeddings for the document + """ class DocOwner(BaseModel): diff --git a/agents-api/agents_api/models/docs/get_doc.py b/agents-api/agents_api/models/docs/get_doc.py index 84cd181ec..4b3996c27 100644 --- a/agents-api/agents_api/models/docs/get_doc.py +++ b/agents-api/agents_api/models/docs/get_doc.py @@ -37,6 +37,7 @@ one=True, transform=lambda d: { "content": [s[1] for s in sorted(d["snippet_data"], key=lambda x: x[0])], + "embeddings": [s[2] for s in sorted(d["snippet_data"], key=lambda x: x[0])], **d, }, ) @@ -68,8 +69,9 @@ def get_doc( doc_id, index, content, + embedding, }, - snippet_data = [index, content] + snippet_data = [index, content, embedding] ?[ id, diff --git a/agents-api/agents_api/models/docs/list_docs.py b/agents-api/agents_api/models/docs/list_docs.py index afdf06c2d..3c095c2db 100644 --- a/agents-api/agents_api/models/docs/list_docs.py +++ b/agents-api/agents_api/models/docs/list_docs.py @@ -34,6 +34,7 @@ Doc, transform=lambda d: { "content": [s[1] for s in sorted(d["snippet_data"], key=lambda x: x[0])], + "embeddings": [s[2] for s in sorted(d["snippet_data"], key=lambda x: x[0])], **d, }, ) @@ -67,8 +68,9 @@ def list_docs( doc_id: id, index, content, + embedding, }}, - snippet_data = [index, content] + snippet_data = [index, content, embedding] ?[ owner_type, diff --git a/typespec/docs/models.tsp b/typespec/docs/models.tsp index ee283fe02..f01b379c6 100644 --- a/typespec/docs/models.tsp +++ b/typespec/docs/models.tsp @@ -23,6 +23,10 @@ model Doc { /** Contents of the document */ content: string | string[]; + + /** Embeddings for the document */ + @visibility("read") + embeddings?: float32[] | float32[][]; } /** Payload for creating a doc */