Merge remote-tracking branch 'origin/main'

dbgpt/agent/core/memory/gpts/gpts_memory.py

@@ -35,7 +35,7 @@ def __init__(
             message_memory if message_memory is not None else DefaultGptsMessageMemory()
         )
 
-        self.messages_cache: defaultdict = defaultdict(List[GptsMessage])
+        self.messages_cache: defaultdict = defaultdict(list)
         self.channels: defaultdict = defaultdict(Queue)
         self.enable_vis_map: defaultdict = defaultdict(bool)
         self.start_round_map: defaultdict = defaultdict(int)
@@ -374,9 +374,9 @@ async def _messages_to_agents_vis(
                     "receiver": message.receiver,
                     "model": message.model_name,
                     "markdown": view_info,
-                    "resource": message.resource_info
-                    if message.resource_info
-                    else None,
+                    "resource": (
+                        message.resource_info if message.resource_info else None
+                    ),
                 }
             )
         return await vis_client.get(VisAgentMessages.vis_tag()).display(
@@ -427,3 +427,20 @@ async def _messages_to_app_link_vis(
         else:
             param["status"] = Status.COMPLETE.value
         return await vis_client.get(VisAppLink.vis_tag()).display(content=param)
+
+    async def chat_messages(
+        self,
+        conv_id: str,
+    ):
+        """Get chat messages."""
+        while True:
+            queue = self.queue(conv_id)
+            if not queue:
+                break
+            item = await queue.get()
+            if item == "[DONE]":
+                queue.task_done()
+                break
+            else:
+                yield item
+                await asyncio.sleep(0.005)

dbgpt/agent/core/plan/awel/agent_operator.py

@@ -22,7 +22,6 @@
 
 # TODO: Don't dependent on MixinLLMOperator
 from dbgpt.model.operators.llm_operator import MixinLLMOperator
-from dbgpt.serve.prompt.api.endpoints import get_service
 from dbgpt.util.i18n_utils import _
 
 from .... import ActionOutput
@@ -291,6 +290,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

dbgpt/agent/core/plan/awel/agent_operator_resource.py

@@ -12,15 +12,14 @@
     ResourceCategory,
     register_resource,
 )
-from dbgpt.serve.prompt.api.endpoints import get_service
-
 from ....resource.base import AgentResource, ResourceType
 from ....resource.manage import get_resource_manager
 from ....util.llm.llm import LLMConfig, LLMStrategyType
 from ...agent_manage import get_agent_manager
 
 
 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 [

dbgpt/app/component_configs.py

@@ -61,6 +61,7 @@ def initialize_components(
     # Register serve apps
     register_serve_apps(system_app, CFG, param.port)
     _initialize_operators()
+    _initialize_code_server(system_app)
 
 
 def _initialize_model_cache(system_app: SystemApp, port: int):
@@ -132,6 +133,7 @@ def _initialize_openapi(system_app: SystemApp):
 
 
 def _initialize_operators():
+    from dbgpt.app.operators.code import CodeMapOperator
     from dbgpt.app.operators.converter import StringToInteger
     from dbgpt.app.operators.datasource import (
         HODatasourceExecutorOperator,
@@ -140,3 +142,9 @@ def _initialize_operators():
     from dbgpt.app.operators.llm import HOLLMOperator, HOStreamingLLMOperator
     from dbgpt.app.operators.rag import HOKnowledgeOperator
     from dbgpt.serve.agent.resource.datasource import DatasourceResource
+
+
+def _initialize_code_server(system_app: SystemApp):
+    from dbgpt.util.code.server import initialize_code_server
+
+    initialize_code_server(system_app)

dbgpt/app/knowledge/api.py

@@ -3,8 +3,9 @@
 import shutil
 import tempfile
 from typing import List
+from pathlib import Path
 
-from fastapi import APIRouter, Depends, File, Form, UploadFile
+from fastapi import APIRouter, Depends, File, Form, UploadFile, HTTPException
 
 from dbgpt._private.config import Config
 from dbgpt.app.knowledge.request.request import (
@@ -332,54 +333,70 @@ 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:
-            if not os.path.exists(os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, space_name)):
-                os.makedirs(os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, space_name))
-            # We can not move temp file in windows system when we open file in context of `with`
-            tmp_fd, tmp_path = tempfile.mkstemp(
-                dir=os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, space_name)
-            )
-            with os.fdopen(tmp_fd, "wb") as tmp:
-                tmp.write(await doc_file.read())
-            shutil.move(
-                tmp_path,
-                os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, space_name, doc_file.filename),
-            )
-            request = KnowledgeDocumentRequest()
-            request.doc_name = doc_name
-            request.doc_type = doc_type
-            request.content = os.path.join(
-                KNOWLEDGE_UPLOAD_ROOT_PATH, space_name, doc_file.filename
-            )
-            space_res = knowledge_space_service.get_knowledge_space(
-                KnowledgeSpaceRequest(name=space_name)
-            )
-            if len(space_res) == 0:
-                # create default space
-                if "default" != space_name:
-                    raise Exception(f"you have not create your knowledge space.")
-                knowledge_space_service.create_knowledge_space(
-                    KnowledgeSpaceRequest(
-                        name=space_name,
-                        desc="first db-gpt rag application",
-                        owner="dbgpt",
-                    )
-                )
-            return Result.succ(
-                knowledge_space_service.create_knowledge_document(
-                    space=space_name, request=request
-                )
-            )
-        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")