yangdx
commited on
Commit
·
27c5884
1
Parent(s):
3d0eaea
Fix linting
Browse files- test_lightrag_ollama_chat.py +23 -12
test_lightrag_ollama_chat.py
CHANGED
|
@@ -19,18 +19,23 @@ from datetime import datetime
|
|
| 19 |
from pathlib import Path
|
| 20 |
from enum import Enum, auto
|
| 21 |
|
|
|
|
| 22 |
class ErrorCode(Enum):
|
| 23 |
"""Error codes for MCP errors"""
|
|
|
|
| 24 |
InvalidRequest = auto()
|
| 25 |
InternalError = auto()
|
| 26 |
|
|
|
|
| 27 |
class McpError(Exception):
|
| 28 |
"""Base exception class for MCP errors"""
|
|
|
|
| 29 |
def __init__(self, code: ErrorCode, message: str):
|
| 30 |
self.code = code
|
| 31 |
self.message = message
|
| 32 |
super().__init__(message)
|
| 33 |
|
|
|
|
| 34 |
DEFAULT_CONFIG = {
|
| 35 |
"server": {
|
| 36 |
"host": "localhost",
|
|
@@ -660,13 +665,17 @@ def test_generate_concurrent() -> None:
|
|
| 660 |
try:
|
| 661 |
async with session.post(url, json=data) as response:
|
| 662 |
if response.status != 200:
|
| 663 |
-
error_msg =
|
|
|
|
|
|
|
| 664 |
if OutputControl.is_verbose():
|
| 665 |
print(f"\n{error_msg}")
|
| 666 |
raise McpError(ErrorCode.InternalError, error_msg)
|
| 667 |
result = await response.json()
|
| 668 |
if "error" in result:
|
| 669 |
-
error_msg =
|
|
|
|
|
|
|
| 670 |
if OutputControl.is_verbose():
|
| 671 |
print(f"\n{error_msg}")
|
| 672 |
raise McpError(ErrorCode.InternalError, error_msg)
|
|
@@ -679,21 +688,23 @@ def test_generate_concurrent() -> None:
|
|
| 679 |
|
| 680 |
async def run_concurrent_requests():
|
| 681 |
prompts = ["第一个问题", "第二个问题", "第三个问题", "第四个问题", "第五个问题"]
|
| 682 |
-
|
| 683 |
async with get_session() as session:
|
| 684 |
-
tasks = [
|
|
|
|
|
|
|
| 685 |
results = await asyncio.gather(*tasks, return_exceptions=True)
|
| 686 |
-
|
| 687 |
# 收集成功和失败的结果
|
| 688 |
success_results = []
|
| 689 |
error_messages = []
|
| 690 |
-
|
| 691 |
for i, result in enumerate(results):
|
| 692 |
if isinstance(result, Exception):
|
| 693 |
error_messages.append(f"Request {i+1} failed: {str(result)}")
|
| 694 |
else:
|
| 695 |
-
success_results.append((i+1, result))
|
| 696 |
-
|
| 697 |
# 如果有任何错误,在打印完所有结果后抛出异常
|
| 698 |
if error_messages:
|
| 699 |
# 先打印成功的结果
|
|
@@ -701,14 +712,14 @@ def test_generate_concurrent() -> None:
|
|
| 701 |
if OutputControl.is_verbose():
|
| 702 |
print(f"\nRequest {req_id} succeeded:")
|
| 703 |
print_json_response(result)
|
| 704 |
-
|
| 705 |
# 打印所有错误信息
|
| 706 |
error_summary = "\n".join(error_messages)
|
| 707 |
raise McpError(
|
| 708 |
ErrorCode.InternalError,
|
| 709 |
-
f"Some concurrent requests failed:\n{error_summary}"
|
| 710 |
)
|
| 711 |
-
|
| 712 |
return results
|
| 713 |
|
| 714 |
if OutputControl.is_verbose():
|
|
@@ -721,7 +732,7 @@ def test_generate_concurrent() -> None:
|
|
| 721 |
for i, result in enumerate(results, 1):
|
| 722 |
print(f"\nRequest {i} result:")
|
| 723 |
print_json_response(result)
|
| 724 |
-
except McpError
|
| 725 |
# 错误信息已经在之前打印过了,这里直接抛出
|
| 726 |
raise
|
| 727 |
|
|
|
|
| 19 |
from pathlib import Path
|
| 20 |
from enum import Enum, auto
|
| 21 |
|
| 22 |
+
|
| 23 |
class ErrorCode(Enum):
|
| 24 |
"""Error codes for MCP errors"""
|
| 25 |
+
|
| 26 |
InvalidRequest = auto()
|
| 27 |
InternalError = auto()
|
| 28 |
|
| 29 |
+
|
| 30 |
class McpError(Exception):
|
| 31 |
"""Base exception class for MCP errors"""
|
| 32 |
+
|
| 33 |
def __init__(self, code: ErrorCode, message: str):
|
| 34 |
self.code = code
|
| 35 |
self.message = message
|
| 36 |
super().__init__(message)
|
| 37 |
|
| 38 |
+
|
| 39 |
DEFAULT_CONFIG = {
|
| 40 |
"server": {
|
| 41 |
"host": "localhost",
|
|
|
|
| 665 |
try:
|
| 666 |
async with session.post(url, json=data) as response:
|
| 667 |
if response.status != 200:
|
| 668 |
+
error_msg = (
|
| 669 |
+
f"Request {request_id} failed with status {response.status}"
|
| 670 |
+
)
|
| 671 |
if OutputControl.is_verbose():
|
| 672 |
print(f"\n{error_msg}")
|
| 673 |
raise McpError(ErrorCode.InternalError, error_msg)
|
| 674 |
result = await response.json()
|
| 675 |
if "error" in result:
|
| 676 |
+
error_msg = (
|
| 677 |
+
f"Request {request_id} returned error: {result['error']}"
|
| 678 |
+
)
|
| 679 |
if OutputControl.is_verbose():
|
| 680 |
print(f"\n{error_msg}")
|
| 681 |
raise McpError(ErrorCode.InternalError, error_msg)
|
|
|
|
| 688 |
|
| 689 |
async def run_concurrent_requests():
|
| 690 |
prompts = ["第一个问题", "第二个问题", "第三个问题", "第四个问题", "第五个问题"]
|
| 691 |
+
|
| 692 |
async with get_session() as session:
|
| 693 |
+
tasks = [
|
| 694 |
+
make_request(session, prompt, i + 1) for i, prompt in enumerate(prompts)
|
| 695 |
+
]
|
| 696 |
results = await asyncio.gather(*tasks, return_exceptions=True)
|
| 697 |
+
|
| 698 |
# 收集成功和失败的结果
|
| 699 |
success_results = []
|
| 700 |
error_messages = []
|
| 701 |
+
|
| 702 |
for i, result in enumerate(results):
|
| 703 |
if isinstance(result, Exception):
|
| 704 |
error_messages.append(f"Request {i+1} failed: {str(result)}")
|
| 705 |
else:
|
| 706 |
+
success_results.append((i + 1, result))
|
| 707 |
+
|
| 708 |
# 如果有任何错误,在打印完所有结果后抛出异常
|
| 709 |
if error_messages:
|
| 710 |
# 先打印成功的结果
|
|
|
|
| 712 |
if OutputControl.is_verbose():
|
| 713 |
print(f"\nRequest {req_id} succeeded:")
|
| 714 |
print_json_response(result)
|
| 715 |
+
|
| 716 |
# 打印所有错误信息
|
| 717 |
error_summary = "\n".join(error_messages)
|
| 718 |
raise McpError(
|
| 719 |
ErrorCode.InternalError,
|
| 720 |
+
f"Some concurrent requests failed:\n{error_summary}",
|
| 721 |
)
|
| 722 |
+
|
| 723 |
return results
|
| 724 |
|
| 725 |
if OutputControl.is_verbose():
|
|
|
|
| 732 |
for i, result in enumerate(results, 1):
|
| 733 |
print(f"\nRequest {i} result:")
|
| 734 |
print_json_response(result)
|
| 735 |
+
except McpError:
|
| 736 |
# 错误信息已经在之前打印过了,这里直接抛出
|
| 737 |
raise
|
| 738 |
|