Skip to content

Commit

Permalink
style:fmt
Browse files Browse the repository at this point in the history
  • Loading branch information
Aries-ckt committed Nov 19, 2024
1 parent 0542168 commit 7f9a89f
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 68 deletions.
1 change: 1 addition & 0 deletions dbgpt/agent/core/plan/awel/agent_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,7 @@ async def get_agent(
prompt_template = None
if self.awel_agent.agent_prompt:
from dbgpt.serve.prompt.api.endpoints import get_service

prompt_service = get_service()
prompt_template = prompt_service.get_template(
self.awel_agent.agent_prompt.code
Expand Down
2 changes: 2 additions & 0 deletions dbgpt/agent/core/plan/awel/agent_operator_resource.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
ResourceCategory,
register_resource,
)

from ....resource.base import AgentResource, ResourceType
from ....resource.manage import get_resource_manager
from ....util.llm.llm import LLMConfig, LLMStrategyType
Expand All @@ -20,6 +21,7 @@

def _agent_resource_prompt_values() -> List[OptionValue]:
from dbgpt.serve.prompt.api.endpoints import get_service

prompt_service = get_service()
prompts = prompt_service.get_target_prompt()
return [
Expand Down
132 changes: 67 additions & 65 deletions dbgpt/app/knowledge/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@
import os
import shutil
import tempfile
from typing import List
from pathlib import Path
from typing import List

from fastapi import APIRouter, Depends, File, Form, UploadFile, HTTPException
from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile

from dbgpt._private.config import Config
from dbgpt.app.knowledge.request.request import (
Expand Down Expand Up @@ -333,70 +333,72 @@ def document_delete(space_name: str, query_request: DocumentQueryRequest):

@router.post("/knowledge/{space_name}/document/upload")
async def document_upload(
space_name: str,
doc_name: str = Form(...),
doc_type: str = Form(...),
doc_file: UploadFile = File(...),
space_name: str,
doc_name: str = Form(...),
doc_type: str = Form(...),
doc_file: UploadFile = File(...),
):
print(f"/document/upload params: {space_name}")
try:
if doc_file:
# Sanitize inputs to prevent path traversal
safe_space_name = os.path.basename(space_name)
safe_filename = os.path.basename(doc_file.filename)

# Create absolute paths and verify they are within allowed directory
upload_dir = os.path.abspath(os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, safe_space_name))
target_path = os.path.abspath(os.path.join(upload_dir, safe_filename))

if not os.path.abspath(KNOWLEDGE_UPLOAD_ROOT_PATH) in target_path:
raise HTTPException(status_code=400, detail="Invalid path detected")

if not os.path.exists(upload_dir):
os.makedirs(upload_dir)

# Create temp file
tmp_fd, tmp_path = tempfile.mkstemp(dir=upload_dir)

try:
with os.fdopen(tmp_fd, "wb") as tmp:
tmp.write(await doc_file.read())

shutil.move(tmp_path, target_path)

request = KnowledgeDocumentRequest()
request.doc_name = doc_name
request.doc_type = doc_type
request.content = target_path

space_res = knowledge_space_service.get_knowledge_space(
KnowledgeSpaceRequest(name=safe_space_name)
)
if len(space_res) == 0:
# create default space
if "default" != safe_space_name:
raise Exception(f"you have not create your knowledge space.")
knowledge_space_service.create_knowledge_space(
KnowledgeSpaceRequest(
name=safe_space_name,
desc="first db-gpt rag application",
owner="dbgpt",
)
)
return Result.succ(
knowledge_space_service.create_knowledge_document(
space=safe_space_name, request=request
)
)
except Exception as e:
# Clean up temp file if anything goes wrong
if os.path.exists(tmp_path):
os.unlink(tmp_path)
raise e

return Result.failed(code="E000X", msg=f"doc_file is None")
except Exception as e:
return Result.failed(code="E000X", msg=f"document add error {e}")
print(f"/document/upload params: {space_name}")
try:
if doc_file:
# Sanitize inputs to prevent path traversal
safe_space_name = os.path.basename(space_name)
safe_filename = os.path.basename(doc_file.filename)

# Create absolute paths and verify they are within allowed directory
upload_dir = os.path.abspath(
os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, safe_space_name)
)
target_path = os.path.abspath(os.path.join(upload_dir, safe_filename))

if not os.path.abspath(KNOWLEDGE_UPLOAD_ROOT_PATH) in target_path:
raise HTTPException(status_code=400, detail="Invalid path detected")

if not os.path.exists(upload_dir):
os.makedirs(upload_dir)

# Create temp file
tmp_fd, tmp_path = tempfile.mkstemp(dir=upload_dir)

try:
with os.fdopen(tmp_fd, "wb") as tmp:
tmp.write(await doc_file.read())

shutil.move(tmp_path, target_path)

request = KnowledgeDocumentRequest()
request.doc_name = doc_name
request.doc_type = doc_type
request.content = target_path

space_res = knowledge_space_service.get_knowledge_space(
KnowledgeSpaceRequest(name=safe_space_name)
)
if len(space_res) == 0:
# create default space
if "default" != safe_space_name:
raise Exception(f"you have not create your knowledge space.")
knowledge_space_service.create_knowledge_space(
KnowledgeSpaceRequest(
name=safe_space_name,
desc="first db-gpt rag application",
owner="dbgpt",
)
)
return Result.succ(
knowledge_space_service.create_knowledge_document(
space=safe_space_name, request=request
)
)
except Exception as e:
# Clean up temp file if anything goes wrong
if os.path.exists(tmp_path):
os.unlink(tmp_path)
raise e

return Result.failed(code="E000X", msg=f"doc_file is None")
except Exception as e:
return Result.failed(code="E000X", msg=f"document add error {e}")


@router.post("/knowledge/{space_name}/document/sync")
Expand Down
4 changes: 1 addition & 3 deletions dbgpt/storage/vector_store/milvus_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -457,9 +457,7 @@ def similar_search_with_scores(
self.vector_field = x.name
# convert to milvus expr filter.
milvus_filter_expr = self.convert_metadata_filters(filters) if filters else None
_, docs_and_scores = self._search(
query=text, k=topk, expr=milvus_filter_expr
)
_, docs_and_scores = self._search(query=text, k=topk, expr=milvus_filter_expr)
if any(score < 0.0 or score > 1.0 for _, score, id in docs_and_scores):
logger.warning(
"similarity score need between" f" 0 and 1, got {docs_and_scores}"
Expand Down

0 comments on commit 7f9a89f

Please sign in to comment.