forked from open-webui/open-webui
collection query fix
This commit is contained in:
parent
1846c1e80d
commit
bc3dd34d8b
1 changed files with 10 additions and 8 deletions
|
@ -29,11 +29,13 @@ from langchain_community.document_loaders import (
|
||||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
import uuid
|
import uuid
|
||||||
|
|
||||||
|
|
||||||
from utils.misc import calculate_sha256, calculate_sha256_string
|
from utils.misc import calculate_sha256, calculate_sha256_string
|
||||||
from utils.utils import get_current_user, get_admin_user
|
from utils.utils import get_current_user, get_admin_user
|
||||||
from config import UPLOAD_DIR, SENTENCE_TRANSFORMER_EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP
|
from config import UPLOAD_DIR, SENTENCE_TRANSFORMER_EMBED_MODEL, CHROMA_CLIENT, CHUNK_SIZE, CHUNK_OVERLAP
|
||||||
|
@ -113,7 +115,7 @@ def query_doc(
|
||||||
# if you use docker use the model from the environment variable
|
# if you use docker use the model from the environment variable
|
||||||
collection = CHROMA_CLIENT.get_collection(
|
collection = CHROMA_CLIENT.get_collection(
|
||||||
name=form_data.collection_name,
|
name=form_data.collection_name,
|
||||||
embedding_function=sentence_transformer_ef
|
embedding_function=sentence_transformer_ef,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# for local development use the default model
|
# for local development use the default model
|
||||||
|
@ -193,13 +195,13 @@ def query_collection(
|
||||||
if 'DOCKER_SENTENCE_TRANSFORMER_EMBED_MODEL' in os.environ:
|
if 'DOCKER_SENTENCE_TRANSFORMER_EMBED_MODEL' in os.environ:
|
||||||
# if you use docker use the model from the environment variable
|
# if you use docker use the model from the environment variable
|
||||||
collection = CHROMA_CLIENT.get_collection(
|
collection = CHROMA_CLIENT.get_collection(
|
||||||
name=form_data.collection_name,
|
name=collection_name,
|
||||||
embedding_function=sentence_transformer_ef
|
embedding_function=sentence_transformer_ef,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
# for local development use the default model
|
# for local development use the default model
|
||||||
collection = CHROMA_CLIENT.get_collection(
|
collection = CHROMA_CLIENT.get_collection(
|
||||||
name=form_data.collection_name,
|
name=collection_name,
|
||||||
)
|
)
|
||||||
|
|
||||||
result = collection.query(
|
result = collection.query(
|
||||||
|
|
Loading…
Reference in a new issue