From 3c7d1cfe98358841753e9902cbf4c524d4616217 Mon Sep 17 00:00:00 2001
From: Shreyas Pimpalgaonkar <shreyas.gp.7@gmail.com>
Date: Wed, 21 Aug 2024 14:58:12 -0400
Subject: [PATCH] rename kgenrichmentresponse (#914)

---
 py/core/base/api/models/__init__.py              |  8 ++++----
 py/core/base/api/models/restructure/responses.py |  4 ++--
 py/sdk/models.py                                 | 13 +------------
 py/sdk/restructure.py                            |  6 +++---
 4 files changed, 10 insertions(+), 21 deletions(-)

diff --git a/py/core/base/api/models/__init__.py b/py/core/base/api/models/__init__.py
index 1bbf8549c..c5506d172 100644
--- a/py/core/base/api/models/__init__.py
+++ b/py/core/base/api/models/__init__.py
@@ -34,8 +34,8 @@
     WrappedUserOverviewResponse,
 )
 from .restructure.responses import (
-    KGEnrichementResponse,
-    WrappedKGEnrichementResponse,
+    KGEnrichmentResponse,
+    WrappedKGEnrichmentResponse,
 )
 from .retrieval.responses import (
     RAGAgentResponse,
@@ -57,8 +57,8 @@
     "IngestionResponse",
     "WrappedIngestionResponse",
     # Restructure Responses
-    "KGEnrichementResponse",
-    "WrappedKGEnrichementResponse",
+    "KGEnrichmentResponse",
+    "WrappedKGEnrichmentResponse",
     # Management Responses
     "PromptResponse",
     "ServerStats",
diff --git a/py/core/base/api/models/restructure/responses.py b/py/core/base/api/models/restructure/responses.py
index e9e62e85e..ed35bd437 100644
--- a/py/core/base/api/models/restructure/responses.py
+++ b/py/core/base/api/models/restructure/responses.py
@@ -6,8 +6,8 @@
 from pydantic import BaseModel
 
 
-class KGEnrichementResponse(BaseModel):
+class KGEnrichmentResponse(BaseModel):
     enriched_content: Dict[str, Any]
 
 
-WrappedKGEnrichementResponse = ResultsWrapper[KGEnrichementResponse]
+WrappedKGEnrichmentResponse = ResultsWrapper[KGEnrichmentResponse]
diff --git a/py/sdk/models.py b/py/sdk/models.py
index 3986a816e..22ded990f 100644
--- a/py/sdk/models.py
+++ b/py/sdk/models.py
@@ -88,17 +88,6 @@ class KGSearchSettings(BaseModel):
     }
 
 
-class KGEnrichmentSettings(BaseModel):
-    leiden_params: dict = Field(
-        default_factory=dict,
-        description="Parameters for the Leiden algorithm.",
-    )
-    generation_config: GenerationConfig = Field(
-        default_factory=GenerationConfig,
-        description="Configuration for text generation during graph enrichment.",
-    )
-
-
 class ProviderConfig(BaseModel, ABC):
     """A base provider configuration class"""
 
@@ -263,7 +252,7 @@ class KGEnrichmentSettings(BaseModel):
         description="The parameters for the Leiden algorithm.",
     )
 
-class KGEnrichementResponse(BaseModel):
+class KGEnrichmentResponse(BaseModel):
     enriched_content: Dict[str, Any]
 
 class UserResponse(BaseModel):
diff --git a/py/sdk/restructure.py b/py/sdk/restructure.py
index 10acb1876..ca93edc2b 100644
--- a/py/sdk/restructure.py
+++ b/py/sdk/restructure.py
@@ -1,18 +1,18 @@
 from typing import Any, Dict, List, Optional, Union
 
 from core.base import Document
-from .models import KGEnrichmentSettings, KGEnrichementResponse
+from .models import KGEnrichmentSettings, KGEnrichmentResponse
 
 class RestructureMethods:
     @staticmethod
     async def enrich_graph(
         client, KGEnrichmentSettings: KGEnrichmentSettings = KGEnrichmentSettings()
-    ) -> KGEnrichementResponse:
+    ) -> KGEnrichmentResponse:
         """
         Perform graph enrichment over the entire graph.
 
         Returns:
-            KGEnrichementResponse: Results of the graph enrichment process.
+            KGEnrichmentResponse: Results of the graph enrichment process.
         """
         if not isinstance(KGEnrichmentSettings, dict):
             KGEnrichmentSettings = KGEnrichmentSettings.model_dump()