58 lines
2.5 KiB
Python
58 lines
2.5 KiB
Python
from __future__ import annotations
|
||
|
||
from .shared import *
|
||
|
||
class LLMGenerateEntryMixin:
|
||
async def generate(self, prompt: str, **kwargs: Any) -> str:
|
||
call_kwargs = dict(kwargs)
|
||
db = call_kwargs.pop("db", None) or call_kwargs.pop("db_session", None)
|
||
if db is None:
|
||
async with AsyncSessionLocal() as session:
|
||
return await self._generate_with_db(session, prompt, **call_kwargs)
|
||
return await self._generate_with_db(db, prompt, **call_kwargs)
|
||
|
||
async def stream_generate(self, prompt: str, **kwargs: Any) -> AsyncIterator[str]:
|
||
call_kwargs = dict(kwargs)
|
||
db = call_kwargs.pop("db", None) or call_kwargs.pop("db_session", None)
|
||
if db is None:
|
||
async with AsyncSessionLocal() as session:
|
||
async for chunk in self._stream_generate_with_db(session, prompt, **call_kwargs):
|
||
yield chunk
|
||
return
|
||
async for chunk in self._stream_generate_with_db(db, prompt, **call_kwargs):
|
||
yield chunk
|
||
|
||
async def generate_structured_object(self, prompt: str, **kwargs: Any) -> StructuredExtractionResult:
|
||
call_kwargs = dict(kwargs)
|
||
db = call_kwargs.pop("db", None) or call_kwargs.pop("db_session", None)
|
||
if db is None:
|
||
async with AsyncSessionLocal() as session:
|
||
return await self._generate_structured_with_db(session, prompt, **call_kwargs)
|
||
return await self._generate_structured_with_db(db, prompt, **call_kwargs)
|
||
|
||
async def _inject_model_context_bundle_prompt(
|
||
self,
|
||
*,
|
||
db: Any,
|
||
tenant_id: int,
|
||
function_name: str,
|
||
prompt: str,
|
||
) -> str:
|
||
raw_prompt = str(prompt or "")
|
||
normalized_function = str(function_name or "").strip().lower()
|
||
if not ModelContextBundleService.should_inject_context(normalized_function):
|
||
return raw_prompt
|
||
# 单测里可能传入 object()/mock 作为 db 占位，此时直接跳过注入。
|
||
if not hasattr(db, "execute"):
|
||
return raw_prompt
|
||
try:
|
||
return await ModelContextBundleService.inject_prompt_with_bundle(
|
||
db,
|
||
tenant_id=int(tenant_id),
|
||
function_name=normalized_function,
|
||
prompt=raw_prompt,
|
||
)
|
||
except Exception as exc:
|
||
_runtime_log("LLM上下文", f"上下文注入失败，已降级原始 prompt：function={normalized_function}, error={exc}")
|
||
return raw_prompt
|