From f62c1b5ccf75133dda39d4fcb5bb11ed76fa9157 Mon Sep 17 00:00:00 2001
From: emrgnt-cmplxty <owen@algofi.org>
Date: Mon, 19 Aug 2024 13:43:17 -0700
Subject: [PATCH] tweak 2

---
 r2r/main/api/routes/ingestion/base.py    | 21 ------------
 r2r/main/api/routes/restructure/base.py  | 20 +++++------
 r2r/main/services/restructure_service.py | 43 +-----------------------
 3 files changed, 9 insertions(+), 75 deletions(-)

diff --git a/r2r/main/api/routes/ingestion/base.py b/r2r/main/api/routes/ingestion/base.py
index c2082f264..e49e37d2c 100644
--- a/r2r/main/api/routes/ingestion/base.py
+++ b/r2r/main/api/routes/ingestion/base.py
@@ -6,7 +6,6 @@
 
 import yaml
 from fastapi import Depends, File, Form, UploadFile
-from fastapi.openapi.models import Example
 
 from r2r.base import ChunkingConfig, R2RException
 from r2r.base.api.models.ingestion.responses import WrappedIngestionResponse
@@ -31,26 +30,6 @@ def load_openapi_extras(self):
             yaml_content = yaml.safe_load(yaml_file)
         return yaml_content
 
-        # paths = yaml_content.get("paths", {})
-
-        # def extract_descriptions(endpoint):
-        #     params = endpoint.get("requestBody", {}).get("content", {}).get("multipart/form-data", {}).get("schema", {}).get("properties", {})
-        #     return {k: v.get("description", "") for k, v in params.items()}
-
-        # ingest_files = paths.get("/v1/ingest_files", {})
-        # # update_files = paths.get("/v1/update_files", {}).get("post", {})
-
-        # return {
-        #     "ingest_files": {
-        #         # "descriptions": extract_descriptions(ingest_files),
-        #         # "operation": ingest_files
-        #     },
-        #     # "update_files": {
-        #     #     "descriptions": extract_descriptions(update_files),
-        #     #     "operation": update_files
-        #     # }
-        # }
-
     def setup_routes(self):
         # Note, we use the following verbose input parameters because FastAPI struggles to handle `File` input and `Body` inputs
         # at the same time. Therefore, we must ues `Form` inputs for the metadata, document_ids, and versions inputs.
diff --git a/r2r/main/api/routes/restructure/base.py b/r2r/main/api/routes/restructure/base.py
index 9660c21de..747d5ed7e 100644
--- a/r2r/main/api/routes/restructure/base.py
+++ b/r2r/main/api/routes/restructure/base.py
@@ -1,6 +1,6 @@
 from typing import List, Optional
 
-from fastapi import Depends
+from fastapi import Depends, Form
 
 from r2r.main.api.routes.base_router import BaseRouter
 from r2r.main.engine import R2REngine
@@ -15,20 +15,16 @@ def setup_routes(self):
         @self.router.post("/enrich_graph")
         @self.base_endpoint
         async def enrich_graph(
-            query: str,
-            entity_types: Optional[List[str]] = None,
-            relationships: Optional[List[str]] = None,
-            generation_config: Optional[dict] = None,
             auth_user=(
                 Depends(self.engine.providers.auth.auth_wrapper)
                 if self.engine.providers.auth
                 else None
             ),
         ):
-            request = {
-                "query": query,
-                "entity_types": entity_types,
-                "relationships": relationships,
-                "generation_config": generation_config,
-            }
-            return await self.engine.enrich_graph(request)
+            """
+            Perform graph enrichment, e.g. GraphRAG, over the ingested documents.
+
+            Returns:
+                Dict[str, Any]: Results of the graph enrichment process.
+            """
+            return await self.engine.aenrich_graph()
diff --git a/r2r/main/services/restructure_service.py b/r2r/main/services/restructure_service.py
index b92d07fcf..ab98182da 100644
--- a/r2r/main/services/restructure_service.py
+++ b/r2r/main/services/restructure_service.py
@@ -30,15 +30,10 @@ def __init__(
             logging_connection,
         )
 
-    async def enrich_graph(
-        self, documents: List[Document] = None
-    ) -> Dict[str, Any]:
+    async def enrich_graph(self) -> Dict[str, Any]:
         """
         Perform graph enrichment on the given documents.
 
-        Args:
-            documents (List[Document]): List of documents to enrich.
-
         Returns:
             Dict[str, Any]: Results of the graph enrichment process.
         """
@@ -54,39 +49,3 @@ async def enrich_graph(
             raise R2RException(
                 status_code=500, message=f"Graph enrichment failed: {str(e)}"
             )
-
-    async def query_graph(self, query: str) -> Dict[str, Any]:
-        """
-        Query the knowledge graph.
-
-        Args:
-            query (str): The query to run against the knowledge graph.
-
-        Returns:
-            Dict[str, Any]: Results of the graph query.
-        """
-        try:
-            results = self.providers.database.graph.query(query)
-            return {"results": results}
-        except Exception as e:
-            logger.error(f"Error querying graph: {str(e)}")
-            raise R2RException(
-                status_code=500, message=f"Graph query failed: {str(e)}"
-            )
-
-    async def get_graph_statistics(self) -> Dict[str, Any]:
-        """
-        Get statistics about the knowledge graph.
-
-        Returns:
-            Dict[str, Any]: Statistics about the knowledge graph.
-        """
-        try:
-            stats = self.providers.database.graph.get_statistics()
-            return stats
-        except Exception as e:
-            logger.error(f"Error getting graph statistics: {str(e)}")
-            raise R2RException(
-                status_code=500,
-                message=f"Failed to retrieve graph statistics: {str(e)}",
-            )