feat(chat): add configurable streaming output
This commit is contained in:
@@ -212,6 +212,7 @@ def save_agent_runtime_config(
|
||||
model.presence_penalty = payload.presencePenalty
|
||||
model.frequency_penalty = payload.frequencyPenalty
|
||||
model.max_token = payload.maxToken
|
||||
model.stream_enabled = payload.streamEnabled
|
||||
db.add(model)
|
||||
OperationLogService.write(
|
||||
db,
|
||||
@@ -436,6 +437,7 @@ def _agent_runtime_config_dict(model: ModelConfig | None) -> dict:
|
||||
"presencePenalty": float(model.presence_penalty) if model is not None and model.presence_penalty is not None else None,
|
||||
"frequencyPenalty": float(model.frequency_penalty) if model is not None and model.frequency_penalty is not None else None,
|
||||
"maxToken": model.max_token if model is not None and model.max_token is not None else 8192,
|
||||
"streamEnabled": model.stream_enabled if model is not None else 1,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -83,7 +83,15 @@ def completions(
|
||||
db: Session = Depends(get_db),
|
||||
current_user: User = Depends(get_current_user),
|
||||
) -> StreamingResponse:
|
||||
return StreamingResponse(_chat_stream(payload, db, current_user), media_type="text/event-stream")
|
||||
return StreamingResponse(
|
||||
_chat_stream(payload, db, current_user),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache, no-transform",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/stop")
|
||||
|
||||
@@ -98,6 +98,7 @@ class AgentRuntimeConfigSaveRequest(BaseModel):
|
||||
presencePenalty: float | None = Field(default=None, ge=-2, le=2)
|
||||
frequencyPenalty: float | None = Field(default=None, ge=-2, le=2)
|
||||
maxToken: int = Field(default=8192, ge=256, le=100000)
|
||||
streamEnabled: int = Field(default=1, ge=0, le=1)
|
||||
|
||||
|
||||
class ModelSaveRequest(BaseModel):
|
||||
|
||||
@@ -60,7 +60,7 @@ class ModelStreamService:
|
||||
model_id=model.id if model is not None else None,
|
||||
model_name=model_name,
|
||||
input_token=_rough_token_count(rag_result.prompt),
|
||||
chunks=_chunk_text(_mock_answer(rag_result)),
|
||||
chunks=_display_chunks(model, _mock_answer(rag_result)),
|
||||
)
|
||||
|
||||
if not rag_result.is_hit and not rag_result.allow_general_knowledge:
|
||||
@@ -68,7 +68,7 @@ class ModelStreamService:
|
||||
model_id=model.id if model is not None else None,
|
||||
model_name=model.model_name if model is not None else "no-hit",
|
||||
input_token=_rough_token_count(rag_result.prompt),
|
||||
chunks=_chunk_text(NO_HIT_ANSWER),
|
||||
chunks=_display_chunks(model, NO_HIT_ANSWER),
|
||||
)
|
||||
|
||||
if model is None:
|
||||
@@ -94,7 +94,7 @@ class ModelStreamService:
|
||||
model_id=model.id if model is not None else None,
|
||||
model_name=model_name,
|
||||
input_token=_rough_token_count(rag_result.prompt),
|
||||
chunks=_async_chunk_text(_mock_answer(rag_result)),
|
||||
chunks=_async_display_chunks(model, _mock_answer(rag_result)),
|
||||
)
|
||||
|
||||
if not rag_result.is_hit and not rag_result.allow_general_knowledge:
|
||||
@@ -102,7 +102,7 @@ class ModelStreamService:
|
||||
model_id=model.id if model is not None else None,
|
||||
model_name=model.model_name if model is not None else "no-hit",
|
||||
input_token=_rough_token_count(rag_result.prompt),
|
||||
chunks=_async_chunk_text(NO_HIT_ANSWER),
|
||||
chunks=_async_display_chunks(model, NO_HIT_ANSWER),
|
||||
)
|
||||
|
||||
if model is None:
|
||||
@@ -130,12 +130,12 @@ def _get_enabled_model(db: Session) -> ModelConfig | None:
|
||||
def _stream_configured_model(model: ModelConfig, rag_result: RagResult) -> Iterator[str]:
|
||||
api_type = model.api_type or "openai_compatible"
|
||||
if model.stream_enabled != 1:
|
||||
return _chunk_text(_call_configured_model(model, rag_result, allow_no_hit=rag_result.allow_general_knowledge))
|
||||
return iter((_call_configured_model(model, rag_result, allow_no_hit=rag_result.allow_general_knowledge),))
|
||||
if api_type == "anthropic_messages":
|
||||
return _stream_anthropic_messages(model, rag_result)
|
||||
if api_type == "openai_compatible":
|
||||
return _stream_openai_compatible_model(model, rag_result)
|
||||
return _chunk_text(_call_configured_model(model, rag_result, allow_no_hit=rag_result.allow_general_knowledge))
|
||||
return iter((_call_configured_model(model, rag_result, allow_no_hit=rag_result.allow_general_knowledge),))
|
||||
|
||||
|
||||
async def _stream_configured_model_async(model: ModelConfig, rag_result: RagResult) -> AsyncIterator[str]:
|
||||
@@ -144,8 +144,7 @@ async def _stream_configured_model_async(model: ModelConfig, rag_result: RagResu
|
||||
answer = await asyncio.to_thread(
|
||||
_call_configured_model, model, rag_result, allow_no_hit=rag_result.allow_general_knowledge
|
||||
)
|
||||
async for chunk in _async_chunk_text(answer):
|
||||
yield chunk
|
||||
yield answer
|
||||
return
|
||||
|
||||
if api_type == "anthropic_messages":
|
||||
@@ -161,8 +160,7 @@ async def _stream_configured_model_async(model: ModelConfig, rag_result: RagResu
|
||||
answer = await asyncio.to_thread(
|
||||
_call_configured_model, model, rag_result, allow_no_hit=rag_result.allow_general_knowledge
|
||||
)
|
||||
async for chunk in _async_chunk_text(answer):
|
||||
yield chunk
|
||||
yield answer
|
||||
|
||||
|
||||
def _stream_openai_compatible_model(model: ModelConfig, rag_result: RagResult) -> Iterator[str]:
|
||||
@@ -479,6 +477,19 @@ def _chunk_text(text: str, *, chunk_size: int = 12) -> Iterator[str]:
|
||||
async def _async_chunk_text(text: str, *, chunk_size: int = 12) -> AsyncIterator[str]:
|
||||
for index in range(0, len(text), chunk_size):
|
||||
yield text[index : index + chunk_size]
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
def _display_chunks(model: ModelConfig | None, text: str) -> Iterator[str]:
|
||||
return _chunk_text(text) if model is None or model.stream_enabled == 1 else iter((text,))
|
||||
|
||||
|
||||
async def _async_display_chunks(model: ModelConfig | None, text: str) -> AsyncIterator[str]:
|
||||
if model is not None and model.stream_enabled != 1:
|
||||
yield text
|
||||
return
|
||||
async for chunk in _async_chunk_text(text):
|
||||
yield chunk
|
||||
|
||||
|
||||
def _string_or_empty(value: Any) -> str:
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import asyncio
|
||||
from decimal import Decimal
|
||||
from unittest.mock import patch
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -9,7 +11,7 @@ from app.models import Base
|
||||
from app.models.admin import Admin
|
||||
from app.models.ai_config import ModelConfig
|
||||
from app.schemas.admin import AgentRuntimeConfigSaveRequest
|
||||
from app.services.model_stream_service import _openai_stream_payload
|
||||
from app.services.model_stream_service import _openai_stream_payload, _stream_configured_model_async
|
||||
from app.services.model_service import ModelClientService, _max_output_tokens
|
||||
from app.services.rag_service import RagResult
|
||||
|
||||
@@ -58,6 +60,7 @@ def test_runtime_config_defaults_to_long_answer_safe_max_tokens():
|
||||
assert result["modelId"] == model.id
|
||||
assert result["modelName"] == "正式模型"
|
||||
assert result["maxToken"] == 8192
|
||||
assert result["streamEnabled"] == 1
|
||||
assert _max_output_tokens(model) == 8192
|
||||
|
||||
|
||||
@@ -76,6 +79,7 @@ def test_saved_runtime_config_is_persisted_on_enabled_model():
|
||||
presencePenalty=0.2,
|
||||
frequencyPenalty=0.4,
|
||||
maxToken=12000,
|
||||
streamEnabled=0,
|
||||
),
|
||||
db=db,
|
||||
current_admin=admin,
|
||||
@@ -89,6 +93,8 @@ def test_saved_runtime_config_is_persisted_on_enabled_model():
|
||||
assert model.top_k == 40
|
||||
assert model.presence_penalty == Decimal("0.20")
|
||||
assert model.frequency_penalty == Decimal("0.40")
|
||||
assert model.stream_enabled == 0
|
||||
assert result["streamEnabled"] == 0
|
||||
|
||||
payload = _openai_stream_payload(
|
||||
model,
|
||||
@@ -101,6 +107,27 @@ def test_saved_runtime_config_is_persisted_on_enabled_model():
|
||||
assert payload["frequency_penalty"] == 0.4
|
||||
|
||||
|
||||
def test_disabled_stream_returns_one_complete_chunk():
|
||||
model = _model()
|
||||
model.stream_enabled = 0
|
||||
answer = "完整回答" * 20
|
||||
rag_result = RagResult(
|
||||
question="测试",
|
||||
knowledge_scopes=[],
|
||||
chunks=[],
|
||||
prompt="测试",
|
||||
allow_general_knowledge=True,
|
||||
)
|
||||
|
||||
async def collect():
|
||||
return [item async for item in _stream_configured_model_async(model, rag_result)]
|
||||
|
||||
with patch("app.services.model_stream_service._call_configured_model", return_value=answer):
|
||||
chunks = asyncio.run(collect())
|
||||
|
||||
assert chunks == [answer]
|
||||
|
||||
|
||||
def test_agent_debug_does_not_truncate_long_answer(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"app.services.model_service._call_configured_model",
|
||||
|
||||
@@ -113,6 +113,15 @@ def test_queued_chat_reports_position_and_completes(monkeypatch):
|
||||
assert released is True
|
||||
|
||||
|
||||
def test_chat_streaming_response_disables_proxy_buffering():
|
||||
payload = type("Payload", (), {"sessionId": 1, "message": "问题"})()
|
||||
response = chat.completions(payload, object(), object())
|
||||
|
||||
assert response.headers["cache-control"] == "no-cache, no-transform"
|
||||
assert response.headers["x-accel-buffering"] == "no"
|
||||
assert response.media_type == "text/event-stream"
|
||||
|
||||
|
||||
def test_sensitive_system_config_never_returns_plaintext():
|
||||
from app.api.admin_settings import _config_dict
|
||||
|
||||
|
||||
Reference in New Issue
Block a user