diff --git a/app/main.py b/app/main.py index 83bb9b6b..2b4bd77f 100644 --- a/app/main.py +++ b/app/main.py @@ -20,7 +20,7 @@ from app.status import StatusManager from app.tools import MapQuestionToSchemaException -from app.py_schemas.schemas import NaturalLanguageQuery, NaturalLanguageQueryResponse, GSQLQueryInfo, BatchDocumentIngest, S3BatchDocumentIngest, SupportAIQuestion +from app.py_schemas.schemas import * from app.log import req_id_cv from app.supportai.retrievers import HNSWOverlapRetriever, HNSWRetriever @@ -280,11 +280,14 @@ def ingestion_status(graphname, status_id: str): else: return {"status": "not found"} -@app.get("/{graphname}/supportai/createvdb/{index_name}") -def create_vdb(graphname, index_name, conn: TigerGraphConnection = Depends(get_db_connection)): - if conn.getVertexCount("HNSWEntrypoint", where='id=="{}"'.format(index_name)) == 0: - res = conn.runInstalledQuery("HNSW_CreateEntrypoint", {"index_name": index_name}) - res = conn.runInstalledQuery("HNSW_BuildIndex", {"index_name": index_name, "v_types": ["DocumentChunk"], "M": 20, "ef_construction": 128}) +@app.post("/{graphname}/supportai/createvdb") +def create_vdb(graphname, config: CreateVectorIndexConfig, conn: TigerGraphConnection = Depends(get_db_connection)): + if conn.getVertexCount("HNSWEntrypoint", where='id=="{}"'.format(config.index_name)) == 0: + res = conn.runInstalledQuery("HNSW_CreateEntrypoint", {"index_name": config.index_name}) + res = conn.runInstalledQuery("HNSW_BuildIndex", {"index_name": config.index_name, + "v_types": config.vertex_types, + "M": config.M, + "ef_construction": config.ef_construction}) return res @app.get("/{graphname}/supportai/deletevdb/{index_name}") diff --git a/app/py_schemas/schemas.py b/app/py_schemas/schemas.py index 9909e23a..3147951d 100644 --- a/app/py_schemas/schemas.py +++ b/app/py_schemas/schemas.py @@ -48,4 +48,10 @@ class Document(BaseModel): document_chunks: List[DocumentChunk] = None entities: List[Dict] = None relationships: List[Dict] = None - document_collection: str = None \ No newline at end of file + document_collection: str = None + +class CreateVectorIndexConfig(BaseModel): + index_name: str + vertex_types: List[str] + M: int = 20 + ef_construction: int = 128 \ No newline at end of file