Merge branch 'main' of https://github.com/tigergraph/natural_language…

…_service

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}")

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
+    document_collection: str = None
+
+class CreateVectorIndexConfig(BaseModel):
+    index_name: str
+    vertex_types: List[str]
+    M: int = 20
+    ef_construction: int = 128