QuivrHQ/quivr
1
1
Fork 0
mirror of https://github.com/QuivrHQ/quivr.git synced 2024-12-15 09:32:22 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix: 🐛 sources

Browse Source 
fixed
This commit is contained in:
Stan Girard 2024-01-28 20:12:28 -08:00
parent 66bb4b1867
commit 67c0c7fb1c
2 changed files with 8 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
backend/llm/knowledge_brain_qa.py
View File

@ -51,19 +51,19 @@ def generate_source(result, brain):
# Iterate over each document # Iterate over each document
for doc in source_documents: for doc in source_documents:
# Check if 'url' is in the document metadata # Check if 'url' is in the document metadata
logger.info(f"Metadata: {doc['metadata']}") logger.info(f"Metadata 1: {doc.metadata}")
is_url = ( is_url = (
"original_file_name" in doc["metadata"] "original_file_name" in doc.metadata
and doc["metadata"]["original_file_name"] is not None and doc.metadata["original_file_name"] is not None
and doc["metadata"]["original_file_name"].startswith("http") and doc.metadata["original_file_name"].startswith("http")
) )
logger.info(f"Is URL: {is_url}") logger.info(f"Is URL: {is_url}")
# Determine the name based on whether it's a URL or a file # Determine the name based on whether it's a URL or a file
name = ( name = (
doc["metadata"]["original_file_name"] doc.metadata["original_file_name"]
if is_url if is_url
else doc["metadata"]["file_name"] else doc.metadata["file_name"]
) )
# Determine the type based on whether it's a URL or a file # Determine the type based on whether it's a URL or a file
@ -71,10 +71,10 @@ def generate_source(result, brain):
# Determine the source URL based on whether it's a URL or a file # Determine the source URL based on whether it's a URL or a file
if is_url: if is_url:
source_url = doc["metadata"]["original_file_name"] source_url = doc.metadata["original_file_name"]
else: else:
source_url = generate_file_signed_url( source_url = generate_file_signed_url(
f"{brain.brain_id}/{doc['metadata']['file_name']}" f"{brain.brain_id}/{doc.metadata['file_name']}"
).get("signedURL", "") ).get("signedURL", "")
# Append a new Sources object to the list # Append a new Sources object to the list

21
backend/llm/tests/test_knowledge_brain_qa.py
View File

@ -8,24 +8,3 @@ def test_generate_source_no_documents():
sources = generate_source(result, brain) sources = generate_source(result, brain)
assert sources == [] assert sources == []
def test_generate_source_with_url():
result = {
"source_documents": [
{
"metadata": {
"original_file_name": "http://example.com",
}
}
]
}
brain = {"brain_id": "123"}
sources = generate_source(result, brain)
assert len(sources) == 1
assert sources[0].name == "http://example.com"
assert sources[0].type == "url"
assert sources[0].source_url == "http://example.com"
assert sources[0].original_file_name == "http://example.com"