14d9540d1b
<!-- .github/pull_request_template.md -->
## Description
- Add support for database deletion when dataset is deleted
- Simplify dataset handler usage in Cognee
## Type of Change
<!-- Please check the relevant option -->
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Code refactoring
- [ ] Performance improvement
- [ ] Other (please specify):
## Screenshots/Videos (if applicable)
<!-- Add screenshots or videos to help explain your changes -->
## Pre-submission Checklist
<!-- Please check all boxes that apply before submitting your PR -->
- [ ] **I have tested my changes thoroughly before submitting this PR**
- [ ] **This PR contains minimal changes necessary to address the
issue/feature**
- [ ] My code follows the project's coding standards and style
guidelines
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have added necessary documentation (if applicable)
- [ ] All new and existing tests pass
- [ ] I have searched existing PRs to ensure this change hasn't been
submitted already
- [ ] I have linked any relevant issues in the description
- [ ] My commits have clear and descriptive messages
## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Bug Fixes**
* Improved dataset deletion: stronger authorization checks and reliable
removal of associated graph and vector storage.
* **Tests**
* Added end-to-end test to verify complete dataset deletion and cleanup
of all related storage components.
<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
from typing import Type, Optional, Coroutine
|
|
from pydantic import BaseModel
|
|
from cognee.infrastructure.llm import get_llm_config
|
|
|
|
|
|
class LLMGateway:
|
|
"""
|
|
Class handles selection of structured output frameworks and LLM functions.
|
|
Class used as a namespace for LLM related functions, should not be instantiated, all methods are static.
|
|
"""
|
|
|
|
@staticmethod
|
|
def acreate_structured_output(
|
|
text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
|
) -> Coroutine:
|
|
llm_config = get_llm_config()
|
|
if llm_config.structured_output_framework.upper() == "BAML":
|
|
from cognee.infrastructure.llm.structured_output_framework.baml.baml_src.extraction import (
|
|
acreate_structured_output,
|
|
)
|
|
|
|
return acreate_structured_output(
|
|
text_input=text_input,
|
|
system_prompt=system_prompt,
|
|
response_model=response_model,
|
|
)
|
|
else:
|
|
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
|
|
get_llm_client,
|
|
)
|
|
|
|
llm_client = get_llm_client()
|
|
return llm_client.acreate_structured_output(
|
|
text_input=text_input,
|
|
system_prompt=system_prompt,
|
|
response_model=response_model,
|
|
**kwargs,
|
|
)
|
|
|
|
@staticmethod
|
|
def create_structured_output(
|
|
text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
|
) -> BaseModel:
|
|
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
|
|
get_llm_client,
|
|
)
|
|
|
|
llm_client = get_llm_client()
|
|
return llm_client.create_structured_output(
|
|
text_input=text_input, system_prompt=system_prompt, response_model=response_model
|
|
)
|
|
|
|
@staticmethod
|
|
def create_transcript(input) -> Coroutine:
|
|
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
|
|
get_llm_client,
|
|
)
|
|
|
|
llm_client = get_llm_client()
|
|
return llm_client.create_transcript(input=input)
|
|
|
|
@staticmethod
|
|
def transcribe_image(input) -> Coroutine:
|
|
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
|
|
get_llm_client,
|
|
)
|
|
|
|
llm_client = get_llm_client()
|
|
return llm_client.transcribe_image(input=input)
|