Skip to content

修复 Dashscope LLM 在 WebSocket 与 HTTP 聊天接口中的流式支持(Fix Dashscope LLM streaming support in WebSocket and HTTP chat endpoints) #404

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+156 −27
Open

修复 Dashscope LLM 在 WebSocket 与 HTTP 聊天接口中的流式支持(Fix Dashscope LLM streaming support in WebSocket and HTTP chat endpoints) #404

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion api/dashscope_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -511,8 +511,23 @@ async def acall(

completion = await self.async_client.chat.completions.create(**api_kwargs)

# For async calls with streaming enabled, wrap the AsyncStream
# into an async generator of plain text chunks so that callers
# can simply `async for text in response`.
if api_kwargs.get("stream", False):
return handle_streaming_response(completion)

async def async_stream_generator():
async for chunk in completion:
log.debug(f"Raw async chunk completion: {chunk}")
try:
parsed_content = parse_stream_response(chunk)
except Exception as e:
log.error(f"Error parsing async stream chunk: {e}")
parsed_content = None
if parsed_content:
yield parsed_content

return async_stream_generator()
else:
return self.parse_chat_completion(completion)
elif model_type == ModelType.EMBEDDER:
Expand Down
92 changes: 77 additions & 15 deletions api/simple_chat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from api.openrouter_client import OpenRouterClient
from api.bedrock_client import BedrockClient
from api.azureai_client import AzureAIClient
from api.dashscope_client import DashscopeClient
from api.rag import RAG
from api.prompts import (
DEEP_RESEARCH_FIRST_ITERATION_PROMPT,
Expand Down Expand Up @@ -63,7 +64,7 @@ class ChatCompletionRequest(BaseModel):
type: Optional[str] = Field("github", description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')")

# model parameters
provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, bedrock, azure)")
provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)")
model: Optional[str] = Field(None, description="Model name for the specified provider")

language: Optional[str] = Field("en", description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')")
Expand Down Expand Up @@ -432,15 +433,31 @@ async def chat_completions_stream(request: ChatCompletionRequest):
model_kwargs=model_kwargs,
model_type=ModelType.LLM
)
elif request.provider == "dashscope":
logger.info(f"Using Dashscope with model: {request.model}")

model = DashscopeClient()
model_kwargs = {
"model": request.model,
"stream": True,
"temperature": model_config["temperature"],
"top_p": model_config["top_p"],
}

api_kwargs = model.convert_inputs_to_api_kwargs(
input=prompt,
model_kwargs=model_kwargs,
model_type=ModelType.LLM,
)
else:
# Initialize Google Generative AI model
# Initialize Google Generative AI model (default provider)
model = genai.GenerativeModel(
model_name=model_config["model"],
generation_config={
"temperature": model_config["temperature"],
"top_p": model_config["top_p"],
"top_k": model_config["top_k"]
}
"top_k": model_config["top_k"],
},
)

# Create a streaming response
Expand Down Expand Up @@ -514,12 +531,29 @@ async def response_stream():
except Exception as e_azure:
logger.error(f"Error with Azure AI API: {str(e_azure)}")
yield f"\nError with Azure AI API: {str(e_azure)}\n\nPlease check that you have set the AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_VERSION environment variables with valid values."
elif request.provider == "dashscope":
try:
logger.info("Making Dashscope API call")
response = await model.acall(
api_kwargs=api_kwargs, model_type=ModelType.LLM
)
# DashscopeClient.acall with stream=True returns an async
# generator of text chunks
async for text in response:
if text:
yield text
except Exception as e_dashscope:
logger.error(f"Error with Dashscope API: {str(e_dashscope)}")
yield (
f"\nError with Dashscope API: {str(e_dashscope)}\n\n"
"Please check that you have set the DASHSCOPE_API_KEY (and optionally "
"DASHSCOPE_WORKSPACE_ID) environment variables with valid values."
)
else:
# Generate streaming response
# Google Generative AI (default provider)
response = model.generate_content(prompt, stream=True)
# Stream the response
for chunk in response:
if hasattr(chunk, 'text'):
if hasattr(chunk, "text"):
yield chunk.text

except Exception as e_outer:
Expand Down Expand Up @@ -648,23 +682,51 @@ async def response_stream():
except Exception as e_fallback:
logger.error(f"Error with Azure AI API fallback: {str(e_fallback)}")
yield f"\nError with Azure AI API fallback: {str(e_fallback)}\n\nPlease check that you have set the AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_VERSION environment variables with valid values."
elif request.provider == "dashscope":
try:
# Create new api_kwargs with the simplified prompt
fallback_api_kwargs = model.convert_inputs_to_api_kwargs(
input=simplified_prompt,
model_kwargs=model_kwargs,
model_type=ModelType.LLM,
)

logger.info("Making fallback Dashscope API call")
fallback_response = await model.acall(
api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM
)

# DashscopeClient.acall (stream=True) returns an async
# generator of text chunks
async for text in fallback_response:
if text:
yield text
except Exception as e_fallback:
logger.error(
f"Error with Dashscope API fallback: {str(e_fallback)}"
)
yield (
f"\nError with Dashscope API fallback: {str(e_fallback)}\n\n"
"Please check that you have set the DASHSCOPE_API_KEY (and optionally "
"DASHSCOPE_WORKSPACE_ID) environment variables with valid values."
)
else:
# Initialize Google Generative AI model
# Google Generative AI fallback (default provider)
model_config = get_model_config(request.provider, request.model)
fallback_model = genai.GenerativeModel(
model_name=model_config["model"],
model_name=model_config["model_kwargs"]["model"],
generation_config={
"temperature": model_config["model_kwargs"].get("temperature", 0.7),
"top_p": model_config["model_kwargs"].get("top_p", 0.8),
"top_k": model_config["model_kwargs"].get("top_k", 40)
}
"top_k": model_config["model_kwargs"].get("top_k", 40),
},
)

# Get streaming response using simplified prompt
fallback_response = fallback_model.generate_content(simplified_prompt, stream=True)
# Stream the fallback response
fallback_response = fallback_model.generate_content(
simplified_prompt, stream=True
)
for chunk in fallback_response:
if hasattr(chunk, 'text'):
if hasattr(chunk, "text"):
yield chunk.text
except Exception as e2:
logger.error(f"Error in fallback streaming response: {str(e2)}")
Expand Down
74 changes: 63 additions & 11 deletions api/websocket_wiki.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -612,14 +612,36 @@ async def handle_websocket_chat(websocket: WebSocket):
await websocket.send_text(error_msg)
# Close the WebSocket connection after sending the error message
await websocket.close()
elif request.provider == "dashscope":
try:
# Get the response and handle it properly using the previously created api_kwargs
logger.info("Making Dashscope API call")
response = await model.acall(
api_kwargs=api_kwargs, model_type=ModelType.LLM
)
# DashscopeClient.acall with stream=True returns an async
# generator of plain text chunks
async for text in response:
if text:
await websocket.send_text(text)
# Explicitly close the WebSocket connection after the response is complete
await websocket.close()
except Exception as e_dashscope:
logger.error(f"Error with Dashscope API: {str(e_dashscope)}")
error_msg = (
f"\nError with Dashscope API: {str(e_dashscope)}\n\n"
"Please check that you have set the DASHSCOPE_API_KEY (and optionally "
"DASHSCOPE_WORKSPACE_ID) environment variables with valid values."
)
await websocket.send_text(error_msg)
# Close the WebSocket connection after sending the error message
await websocket.close()
else:
# Generate streaming response
# Google Generative AI (default provider)
response = model.generate_content(prompt, stream=True)
# Stream the response
for chunk in response:
if hasattr(chunk, 'text'):
await websocket.send_text(chunk.text)
# Explicitly close the WebSocket connection after the response is complete
await websocket.close()

except Exception as e_outer:
Expand Down Expand Up @@ -729,23 +751,52 @@ async def handle_websocket_chat(websocket: WebSocket):
logger.error(f"Error with Azure AI API fallback: {str(e_fallback)}")
error_msg = f"\nError with Azure AI API fallback: {str(e_fallback)}\n\nPlease check that you have set the AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_VERSION environment variables with valid values."
await websocket.send_text(error_msg)
elif request.provider == "dashscope":
try:
# Create new api_kwargs with the simplified prompt
fallback_api_kwargs = model.convert_inputs_to_api_kwargs(
input=simplified_prompt,
model_kwargs=model_kwargs,
model_type=ModelType.LLM,
)

logger.info("Making fallback Dashscope API call")
fallback_response = await model.acall(
api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM
)

# DashscopeClient.acall (stream=True) returns an async
# generator of text chunks
async for text in fallback_response:
if text:
await websocket.send_text(text)
except Exception as e_fallback:
logger.error(
f"Error with Dashscope API fallback: {str(e_fallback)}"
)
error_msg = (
f"\nError with Dashscope API fallback: {str(e_fallback)}\n\n"
"Please check that you have set the DASHSCOPE_API_KEY (and optionally "
"DASHSCOPE_WORKSPACE_ID) environment variables with valid values."
)
await websocket.send_text(error_msg)
else:
# Initialize Google Generative AI model
# Google Generative AI fallback (default provider)
model_config = get_model_config(request.provider, request.model)
fallback_model = genai.GenerativeModel(
model_name=model_config["model"],
model_name=model_config["model_kwargs"]["model"],
generation_config={
"temperature": model_config["model_kwargs"].get("temperature", 0.7),
"top_p": model_config["model_kwargs"].get("top_p", 0.8),
"top_k": model_config["model_kwargs"].get("top_k", 40)
}
"top_k": model_config["model_kwargs"].get("top_k", 40),
},
)

# Get streaming response using simplified prompt
fallback_response = fallback_model.generate_content(simplified_prompt, stream=True)
# Stream the fallback response
fallback_response = fallback_model.generate_content(
simplified_prompt, stream=True
)
for chunk in fallback_response:
if hasattr(chunk, 'text'):
if hasattr(chunk, "text"):
await websocket.send_text(chunk.text)
except Exception as e2:
logger.error(f"Error in fallback streaming response: {str(e2)}")
Expand All @@ -767,3 +818,4 @@ async def handle_websocket_chat(websocket: WebSocket):
await websocket.close()
except:
pass