f75e35c337
<!-- .github/pull_request_template.md --> ## Description <!-- Provide a clear description of the changes in this PR --> ## 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 - **New Features** • Graph visualizations now allow exporting to a user-specified file path for more flexible output management. • The text embedding process has been enhanced with an additional tokenizer option for improved performance. • A new `ExtendableDataPoint` class has been introduced for future extensions. • New JSON files for companies and individuals have been added to facilitate testing and data processing. - **Improvements** • Search functionality now uses updated identifiers for more reliable content retrieval. • Metadata handling has been streamlined across various classes by removing unnecessary type specifications. • Enhanced serialization of properties in the Neo4j adapter for improved handling of complex structures. • The setup process for databases has been improved with a new asynchronous setup function. - **Chores** • Dependency and configuration updates improve overall stability and performance. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
from cognee.infrastructure.databases.vector import get_vector_engine
|
|
from cognee.tasks.completion.exceptions import NoRelevantDataFound
|
|
from cognee.infrastructure.llm.get_llm_client import get_llm_client
|
|
from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt
|
|
|
|
|
|
async def query_completion(query: str) -> list:
|
|
"""
|
|
|
|
Executes a query against a vector database and computes a relevant response using an LLM.
|
|
|
|
Parameters:
|
|
- query (str): The query string to compute.
|
|
|
|
Returns:
|
|
- list: Answer to the query.
|
|
|
|
Notes:
|
|
- Limits the search to the top 1 matching chunk for simplicity and relevance.
|
|
- Ensure that the vector database and LLM client are properly configured and accessible.
|
|
- The response model used for the LLM output is expected to be a string.
|
|
|
|
"""
|
|
vector_engine = get_vector_engine()
|
|
|
|
found_chunks = await vector_engine.search("DocumentChunk_text", query, limit=1)
|
|
|
|
if len(found_chunks) == 0:
|
|
raise NoRelevantDataFound
|
|
|
|
args = {
|
|
"question": query,
|
|
"context": found_chunks[0].payload["text"],
|
|
}
|
|
user_prompt = render_prompt("context_for_question.txt", args)
|
|
system_prompt = read_query_prompt("answer_simple_question.txt")
|
|
|
|
llm_client = get_llm_client()
|
|
computed_answer = await llm_client.acreate_structured_output(
|
|
text_input=user_prompt,
|
|
system_prompt=system_prompt,
|
|
response_model=str,
|
|
)
|
|
|
|
return [computed_answer]
|