feat: add support to pass custom parameters in llm adapters during cognify (#1802)
<!-- .github/pull_request_template.md --> ## Description <!-- Please provide a clear, human-generated description of the changes in this PR. DO NOT use AI-generated descriptions. We want to understand your thought process and reasoning. --> A user raised this issue/feature request: https://github.com/topoteretes/cognee/issues/1784 This PR implements this in a way, but I may have misunderstood the request. If I'm correct, I think it makes sense to support this, if our users need it. ## Type of Change <!-- Please check the relevant option --> - [ ] 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 commit is contained in:
commit
821852f2c3
10 changed files with 26 additions and 16 deletions
|
|
@ -53,6 +53,7 @@ async def cognify(
|
||||||
custom_prompt: Optional[str] = None,
|
custom_prompt: Optional[str] = None,
|
||||||
temporal_cognify: bool = False,
|
temporal_cognify: bool = False,
|
||||||
data_per_batch: int = 20,
|
data_per_batch: int = 20,
|
||||||
|
**kwargs
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Transform ingested data into a structured knowledge graph.
|
Transform ingested data into a structured knowledge graph.
|
||||||
|
|
@ -223,6 +224,7 @@ async def cognify(
|
||||||
config=config,
|
config=config,
|
||||||
custom_prompt=custom_prompt,
|
custom_prompt=custom_prompt,
|
||||||
chunks_per_batch=chunks_per_batch,
|
chunks_per_batch=chunks_per_batch,
|
||||||
|
**kwargs,
|
||||||
)
|
)
|
||||||
|
|
||||||
# By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for
|
# By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for
|
||||||
|
|
@ -251,6 +253,7 @@ async def get_default_tasks( # TODO: Find out a better way to do this (Boris's
|
||||||
config: Config = None,
|
config: Config = None,
|
||||||
custom_prompt: Optional[str] = None,
|
custom_prompt: Optional[str] = None,
|
||||||
chunks_per_batch: int = 100,
|
chunks_per_batch: int = 100,
|
||||||
|
**kwargs,
|
||||||
) -> list[Task]:
|
) -> list[Task]:
|
||||||
if config is None:
|
if config is None:
|
||||||
ontology_config = get_ontology_env_config()
|
ontology_config = get_ontology_env_config()
|
||||||
|
|
@ -288,6 +291,7 @@ async def get_default_tasks( # TODO: Find out a better way to do this (Boris's
|
||||||
config=config,
|
config=config,
|
||||||
custom_prompt=custom_prompt,
|
custom_prompt=custom_prompt,
|
||||||
task_config={"batch_size": chunks_per_batch},
|
task_config={"batch_size": chunks_per_batch},
|
||||||
|
**kwargs,
|
||||||
), # Generate knowledge graphs from the document chunks.
|
), # Generate knowledge graphs from the document chunks.
|
||||||
Task(
|
Task(
|
||||||
summarize_text,
|
summarize_text,
|
||||||
|
|
|
||||||
|
|
@ -11,7 +11,7 @@ class LLMGateway:
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def acreate_structured_output(
|
def acreate_structured_output(
|
||||||
text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||||
) -> Coroutine:
|
) -> Coroutine:
|
||||||
llm_config = get_llm_config()
|
llm_config = get_llm_config()
|
||||||
if llm_config.structured_output_framework.upper() == "BAML":
|
if llm_config.structured_output_framework.upper() == "BAML":
|
||||||
|
|
@ -31,7 +31,7 @@ class LLMGateway:
|
||||||
|
|
||||||
llm_client = get_llm_client()
|
llm_client = get_llm_client()
|
||||||
return llm_client.acreate_structured_output(
|
return llm_client.acreate_structured_output(
|
||||||
text_input=text_input, system_prompt=system_prompt, response_model=response_model
|
text_input=text_input, system_prompt=system_prompt, response_model=response_model, **kwargs
|
||||||
)
|
)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
|
|
||||||
|
|
@ -10,7 +10,7 @@ from cognee.infrastructure.llm.config import (
|
||||||
|
|
||||||
|
|
||||||
async def extract_content_graph(
|
async def extract_content_graph(
|
||||||
content: str, response_model: Type[BaseModel], custom_prompt: Optional[str] = None
|
content: str, response_model: Type[BaseModel], custom_prompt: Optional[str] = None, **kwargs
|
||||||
):
|
):
|
||||||
if custom_prompt:
|
if custom_prompt:
|
||||||
system_prompt = custom_prompt
|
system_prompt = custom_prompt
|
||||||
|
|
@ -30,7 +30,7 @@ async def extract_content_graph(
|
||||||
system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory)
|
system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory)
|
||||||
|
|
||||||
content_graph = await LLMGateway.acreate_structured_output(
|
content_graph = await LLMGateway.acreate_structured_output(
|
||||||
content, system_prompt, response_model
|
content, system_prompt, response_model, **kwargs
|
||||||
)
|
)
|
||||||
|
|
||||||
return content_graph
|
return content_graph
|
||||||
|
|
|
||||||
|
|
@ -52,7 +52,7 @@ class AnthropicAdapter(LLMInterface):
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def acreate_structured_output(
|
async def acreate_structured_output(
|
||||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||||
) -> BaseModel:
|
) -> BaseModel:
|
||||||
"""
|
"""
|
||||||
Generate a response from a user query.
|
Generate a response from a user query.
|
||||||
|
|
|
||||||
|
|
@ -80,7 +80,7 @@ class GeminiAdapter(LLMInterface):
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def acreate_structured_output(
|
async def acreate_structured_output(
|
||||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||||
) -> BaseModel:
|
) -> BaseModel:
|
||||||
"""
|
"""
|
||||||
Generate a response from a user query.
|
Generate a response from a user query.
|
||||||
|
|
|
||||||
|
|
@ -80,7 +80,7 @@ class GenericAPIAdapter(LLMInterface):
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def acreate_structured_output(
|
async def acreate_structured_output(
|
||||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||||
) -> BaseModel:
|
) -> BaseModel:
|
||||||
"""
|
"""
|
||||||
Generate a response from a user query.
|
Generate a response from a user query.
|
||||||
|
|
|
||||||
|
|
@ -69,7 +69,7 @@ class MistralAdapter(LLMInterface):
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def acreate_structured_output(
|
async def acreate_structured_output(
|
||||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||||
) -> BaseModel:
|
) -> BaseModel:
|
||||||
"""
|
"""
|
||||||
Generate a response from the user query.
|
Generate a response from the user query.
|
||||||
|
|
|
||||||
|
|
@ -76,7 +76,7 @@ class OllamaAPIAdapter(LLMInterface):
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def acreate_structured_output(
|
async def acreate_structured_output(
|
||||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||||
) -> BaseModel:
|
) -> BaseModel:
|
||||||
"""
|
"""
|
||||||
Generate a structured output from the LLM using the provided text and system prompt.
|
Generate a structured output from the LLM using the provided text and system prompt.
|
||||||
|
|
@ -123,7 +123,7 @@ class OllamaAPIAdapter(LLMInterface):
|
||||||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def create_transcript(self, input_file: str) -> str:
|
async def create_transcript(self, input_file: str, **kwargs) -> str:
|
||||||
"""
|
"""
|
||||||
Generate an audio transcript from a user query.
|
Generate an audio transcript from a user query.
|
||||||
|
|
||||||
|
|
@ -162,7 +162,7 @@ class OllamaAPIAdapter(LLMInterface):
|
||||||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def transcribe_image(self, input_file: str) -> str:
|
async def transcribe_image(self, input_file: str, **kwargs) -> str:
|
||||||
"""
|
"""
|
||||||
Transcribe content from an image using base64 encoding.
|
Transcribe content from an image using base64 encoding.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -112,7 +112,7 @@ class OpenAIAdapter(LLMInterface):
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def acreate_structured_output(
|
async def acreate_structured_output(
|
||||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||||
) -> BaseModel:
|
) -> BaseModel:
|
||||||
"""
|
"""
|
||||||
Generate a response from a user query.
|
Generate a response from a user query.
|
||||||
|
|
@ -154,6 +154,7 @@ class OpenAIAdapter(LLMInterface):
|
||||||
api_version=self.api_version,
|
api_version=self.api_version,
|
||||||
response_model=response_model,
|
response_model=response_model,
|
||||||
max_retries=self.MAX_RETRIES,
|
max_retries=self.MAX_RETRIES,
|
||||||
|
**kwargs,
|
||||||
)
|
)
|
||||||
except (
|
except (
|
||||||
ContentFilterFinishReasonError,
|
ContentFilterFinishReasonError,
|
||||||
|
|
@ -180,6 +181,7 @@ class OpenAIAdapter(LLMInterface):
|
||||||
# api_base=self.fallback_endpoint,
|
# api_base=self.fallback_endpoint,
|
||||||
response_model=response_model,
|
response_model=response_model,
|
||||||
max_retries=self.MAX_RETRIES,
|
max_retries=self.MAX_RETRIES,
|
||||||
|
**kwargs,
|
||||||
)
|
)
|
||||||
except (
|
except (
|
||||||
ContentFilterFinishReasonError,
|
ContentFilterFinishReasonError,
|
||||||
|
|
@ -205,7 +207,7 @@ class OpenAIAdapter(LLMInterface):
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
def create_structured_output(
|
def create_structured_output(
|
||||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||||
) -> BaseModel:
|
) -> BaseModel:
|
||||||
"""
|
"""
|
||||||
Generate a response from a user query.
|
Generate a response from a user query.
|
||||||
|
|
@ -245,6 +247,7 @@ class OpenAIAdapter(LLMInterface):
|
||||||
api_version=self.api_version,
|
api_version=self.api_version,
|
||||||
response_model=response_model,
|
response_model=response_model,
|
||||||
max_retries=self.MAX_RETRIES,
|
max_retries=self.MAX_RETRIES,
|
||||||
|
**kwargs,
|
||||||
)
|
)
|
||||||
|
|
||||||
@retry(
|
@retry(
|
||||||
|
|
@ -254,7 +257,7 @@ class OpenAIAdapter(LLMInterface):
|
||||||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def create_transcript(self, input):
|
async def create_transcript(self, input, **kwargs):
|
||||||
"""
|
"""
|
||||||
Generate an audio transcript from a user query.
|
Generate an audio transcript from a user query.
|
||||||
|
|
||||||
|
|
@ -281,6 +284,7 @@ class OpenAIAdapter(LLMInterface):
|
||||||
api_base=self.endpoint,
|
api_base=self.endpoint,
|
||||||
api_version=self.api_version,
|
api_version=self.api_version,
|
||||||
max_retries=self.MAX_RETRIES,
|
max_retries=self.MAX_RETRIES,
|
||||||
|
**kwargs,
|
||||||
)
|
)
|
||||||
|
|
||||||
return transcription
|
return transcription
|
||||||
|
|
@ -292,7 +296,7 @@ class OpenAIAdapter(LLMInterface):
|
||||||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||||
reraise=True,
|
reraise=True,
|
||||||
)
|
)
|
||||||
async def transcribe_image(self, input) -> BaseModel:
|
async def transcribe_image(self, input, **kwargs) -> BaseModel:
|
||||||
"""
|
"""
|
||||||
Generate a transcription of an image from a user query.
|
Generate a transcription of an image from a user query.
|
||||||
|
|
||||||
|
|
@ -337,4 +341,5 @@ class OpenAIAdapter(LLMInterface):
|
||||||
api_version=self.api_version,
|
api_version=self.api_version,
|
||||||
max_completion_tokens=300,
|
max_completion_tokens=300,
|
||||||
max_retries=self.MAX_RETRIES,
|
max_retries=self.MAX_RETRIES,
|
||||||
|
**kwargs,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -97,6 +97,7 @@ async def extract_graph_from_data(
|
||||||
graph_model: Type[BaseModel],
|
graph_model: Type[BaseModel],
|
||||||
config: Config = None,
|
config: Config = None,
|
||||||
custom_prompt: Optional[str] = None,
|
custom_prompt: Optional[str] = None,
|
||||||
|
**kwargs,
|
||||||
) -> List[DocumentChunk]:
|
) -> List[DocumentChunk]:
|
||||||
"""
|
"""
|
||||||
Extracts and integrates a knowledge graph from the text content of document chunks using a specified graph model.
|
Extracts and integrates a knowledge graph from the text content of document chunks using a specified graph model.
|
||||||
|
|
@ -111,7 +112,7 @@ async def extract_graph_from_data(
|
||||||
|
|
||||||
chunk_graphs = await asyncio.gather(
|
chunk_graphs = await asyncio.gather(
|
||||||
*[
|
*[
|
||||||
extract_content_graph(chunk.text, graph_model, custom_prompt=custom_prompt)
|
extract_content_graph(chunk.text, graph_model, custom_prompt=custom_prompt, **kwargs)
|
||||||
for chunk in data_chunks
|
for chunk in data_chunks
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue