pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2026-02-05 01:55:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: Add a custom header to the SDK for chatting with the agent. (#12430)

Browse Source 
### What problem does this PR solve?

As title.

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

Co-authored-by: Liu An <asiro@qq.com>
This commit is contained in:
zhanglei
2026-02-03 11:01:18 +08:00
committed by GitHub
parent aa8d0a36f1
commit 7cbe8b5b53
4 changed files with 20 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
common/mcp_tool_call_conn.py
View File

@ -42,9 +42,10 @@ class ToolCallSession(Protocol):
class MCPToolCallSession(ToolCallSession):
_ALL_INSTANCES: weakref.WeakSet["MCPToolCallSession"] = weakref.WeakSet()
def __init__(self, mcp_server: Any, server_variables: dict[str, Any] | None = None) -> None:
def __init__(self, mcp_server: Any, server_variables: dict[str, Any] | None = None, custom_header = None) -> None:
self.__class__._ALL_INSTANCES.add(self)
self._custom_header = custom_header
self._mcp_server = mcp_server
self._server_variables = server_variables or {}
self._queue = asyncio.Queue()
@ -59,6 +60,7 @@ class MCPToolCallSession(ToolCallSession):
async def _mcp_server_loop(self) -> None:
url = self._mcp_server.url.strip()
raw_headers: dict[str, str] = self._mcp_server.headers or {}
custom_header: dict[str, str] = self._custom_header or {}
headers: dict[str, str] = {}
for h, v in raw_headers.items():
@ -67,6 +69,11 @@ class MCPToolCallSession(ToolCallSession):
if nh.strip() and nv.strip().strip("Bearer"):
headers[nh] = nv
for h, v in custom_header.items():
nh = Template(h).safe_substitute(custom_header)
nv = Template(v).safe_substitute(custom_header)
headers[nh] = nv
if self._mcp_server.server_type == MCPServerType.SSE:
# SSE transport
try: