Patch/fix kg issue (#1208)

* revamp agent streaming

* revamp agent streaming

* fix kg bug

py/core/base/providers/llm.py

@@ -18,7 +18,7 @@
 
 class CompletionConfig(ProviderConfig):
     provider: Optional[str] = None
-    generation_config: Optional[GenerationConfig] = None
+    generation_config: GenerationConfig = GenerationConfig()
     concurrent_request_limit: int = 256
     max_retries: int = 2
     initial_backoff: float = 1.0

py/core/main/hatchet/restructure_workflow.py

@@ -56,13 +56,16 @@ async def kg_extract_and_store(self, context: Context) -> None:
                 entity_types=entity_types,
                 relation_types=relation_types,
             )
-
             # Set restructure status to 'success' if completed successfully
             if len(errors) == 0:
                 document_overview.restructuring_status = (
                     RestructureStatus.SUCCESS
                 )
+                await self.restructure_service.providers.database.relational.upsert_documents_overview(
+                    document_overview
+                )
             else:
+
                 document_overview.restructuring_status = (
                     RestructureStatus.FAILURE
                 )
@@ -80,10 +83,12 @@ async def kg_extract_and_store(self, context: Context) -> None:
             await self.restructure_service.providers.database.relational.upsert_documents_overview(
                 document_overview
             )
-            logger.error(
-                f"Error in kg_extract_and_store for document {document_id}: {str(e)}"
+            raise R2RDocumentProcessingError(
+                error_message=e,
+                document_id=document_id,
             )
 
+
         return {"result": None}
 
 

py/core/providers/kg/neo4j/graph_queries.py

@@ -1,29 +1,10 @@
-UNIQUE_CONSTRAINTS = [
-    # "create constraint chunk_id if not exists for (c:__Chunk__) require c.id is unique;",
-    # "create constraint document_id if not exists for (d:__Document__) require d.id is unique;",
-    # "create constraint entity_id if not exists for (c:__Community__) require c.community is unique;",
-    # "create constraint entity_id if not exists for (e:__Entity__) require e.id is unique;",
-    # "create constraint entity_title if not exists for (e:__Entity__) require e.name is unique;",
-    # "create constraint entity_title if not exists for (e:__Covariate__) require e.title is unique;",
-    # "create constraint related_id if not exists for ()-[rel:RELATED]->() require rel.id is unique;"
-]
+UNIQUE_CONSTRAINTS = []
 
 GET_CHUNKS_QUERY = """
 MATCH (c:__Chunk__)
 RETURN c
 """
 
-# class DocumentFragment(BaseModel):
-#     """A fragment extracted from a document."""
-
-#     id: uuid.UUID
-#     type: FragmentType
-#     data: DataType
-#     metadata: dict
-#     document_id: uuid.UUID
-#     extraction_id: uuid.UUID
-
-
 PUT_CHUNKS_QUERY = """
 MERGE (c:__Chunk__ {id:value.id})
 SET c += value {.type, .data, .metadata, .document_id, .extraction_id}

py/core/providers/prompts/defaults/rag_agent.yaml

@@ -1,6 +1,6 @@
 rag_agent:
   template: >
-    ## You are a helpful agent that can search for information.
+    ### You are a helpful agent that can search for information.
 
 
     When asked a question, perform a search to find relevant information and provide a response.

py/pyproject.toml

@@ -120,7 +120,7 @@ profile = "black"
 
 [tool.mypy]
 ignore_missing_imports = true
-exclude = 'playground/.*|deprecated/.*|dump/.*|docs/source|vecs/*'
+exclude = 'playground/.*|deprecated/.*|dump/.*|docs/source|vecs/*|core/examples/*|sdk/examples/*|cli/examples/*|tests/*'
 
 [[tool.mypy.overrides]]
 module = "yaml"