ydy0615/llm-in-text
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enhance LLM functionality with PRO mode support and improved prompt handling

Browse Source 
- Added support for PRO mode in LLM with specific instruction handling and context awareness.
- Updated prompt building functions to include prefill options for better context management.
- Introduced new inline examples for PRO mode in JSON format.
- Enhanced system prompts to reflect PRO mode capabilities and rules.
- Modified API endpoints to accommodate new parameters and ensure backward compatibility.
- Improved test cases to validate new functionality and ensure comprehensive coverage.
This commit is contained in:
“ydy0615”
2026-06-02 21:23:34 +08:00
parent b82c6d392d
commit 2c7a02f587
12 changed files with 191 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/llm.py
+21 -7
View File
@@ -18,6 +18,9 @@ load_dotenv()
LLM_BASE_URL = os.getenv('LLM_BASE_URL', 'http://localhost:11434/v1/')
LLM_API_KEY = os.getenv('LLM_API_KEY', 'ollama')
# Auth headers for upstream LLM service (OpenAI-compatible Bearer token)
LLM_HEADERS = {'Authorization': f'Bearer {LLM_API_KEY}'}
# Model names (backward compat: fall back to OLLAMA_MODEL if LLM_MODEL not set)
_raw_model = os.getenv('LLM_MODEL') or os.getenv('OLLAMA_MODEL', 'gpt-oss:20b')
LLM_MODEL = _raw_model.strip() if _raw_model else 'gpt-oss:20b'
@@ -73,6 +76,7 @@ def _build_chat_payload(
thinking: str | None = None,
model: str | None = None,
use_pro_model: bool = False,
prefill: str | None = None,
) -> dict:
messages = []
sys_prompt = _resolve_system_prompt(system_prompt)
@@ -80,6 +84,9 @@ def _build_chat_payload(
messages.append({'role': 'system', 'content': sys_prompt})
messages.append({'role': 'user', 'content': prompt})
if prefill:
messages.append({'role': 'assistant', 'content': prefill})
payload = {
'model': _resolve_model_name(model, use_pro_model=use_pro_model),
'messages': messages,
@@ -101,6 +108,7 @@ def _build_chat_stream_payload(
thinking: str | None = None,
model: str | None = None,
use_pro_model: bool = False,
prefill: str | None = None,
) -> dict:
messages = []
sys_prompt = _resolve_system_prompt(system_prompt)
@@ -108,6 +116,9 @@ def _build_chat_stream_payload(
messages.append({'role': 'system', 'content': sys_prompt})
messages.append({'role': 'user', 'content': prompt})
if prefill:
messages.append({'role': 'assistant', 'content': prefill})
payload = {
'model': _resolve_model_name(model, use_pro_model=use_pro_model),
'messages': messages,
@@ -145,6 +156,7 @@ async def call_ollama(
thinking: str | None = None,
model: str | None = None,
use_pro_model: bool = False,
prefill: str | None = None,
) -> dict:
"""Call OpenAI-compatible chat completions (non-streaming) and return content/thinking."""
start = time.perf_counter()
@@ -161,13 +173,13 @@ async def call_ollama(
payload = _build_chat_payload(
prompt=prompt, system_prompt=system_prompt, temperature=temperature,
thinking=thinking, model=model, use_pro_model=use_pro_model,
thinking=thinking, model=model, use_pro_model=use_pro_model, prefill=prefill,
)
http_timeout = httpx.Timeout(connect=10.0, read=None, write=30.0, pool=30.0)
try:
async with httpx.AsyncClient(base_url=LLM_BASE_URL, timeout=http_timeout) as client:
async with httpx.AsyncClient(base_url=LLM_BASE_URL, headers=LLM_HEADERS, timeout=http_timeout) as client:
resp = await asyncio.wait_for(
client.post('/chat/completions', json=payload), timeout=COMPLETION_TIMEOUT,
)
@@ -229,6 +241,7 @@ async def stream_ollama(
thinking: str | None = None,
model: str | None = None,
use_pro_model: bool = False,
prefill: str | None = None,
) -> AsyncIterator[str]:
"""Stream text deltas from OpenAI-compatible chat completions."""
start = time.perf_counter()
@@ -246,13 +259,13 @@ async def stream_ollama(
payload = _build_chat_stream_payload(
prompt=prompt, system_prompt=system_prompt, temperature=temperature,
thinking=thinking, model=model, use_pro_model=use_pro_model,
thinking=thinking, model=model, use_pro_model=use_pro_model, prefill=prefill,
)
http_timeout = httpx.Timeout(connect=10.0, read=None, write=30.0, pool=30.0)
try:
async with httpx.AsyncClient(base_url=LLM_BASE_URL, timeout=http_timeout) as client:
async with httpx.AsyncClient(base_url=LLM_BASE_URL, headers=LLM_HEADERS, timeout=http_timeout) as client:
try:
async with client.stream('POST', '/chat/completions', json=payload) as response:
response.raise_for_status()
@@ -352,6 +365,7 @@ async def stream_ollama_events(
model: str | None = None,
use_pro_model: bool = False,
enable_thinking: bool = True,
prefill: str | None = None,
timeout: float | None = None,
) -> AsyncIterator[tuple[Literal['thinking', 'content'], str]]:
"""Stream (event_type, payload) tuples from OpenAI-compatible chat completions."""
@@ -370,7 +384,7 @@ async def stream_ollama_events(
payload = _build_chat_stream_payload(
prompt=prompt, system_prompt=system_prompt, temperature=temperature,
thinking=thinking if enable_thinking else None, model=model, use_pro_model=use_pro_model,
thinking=thinking if enable_thinking else None, model=model, use_pro_model=use_pro_model, prefill=prefill,
)
effective_timeout = timeout if timeout is not None else COMPLETION_TIMEOUT
@@ -378,7 +392,7 @@ async def stream_ollama_events(
sent_thinking = False
try:
async with httpx.AsyncClient(base_url=LLM_BASE_URL, timeout=http_timeout) as client:
async with httpx.AsyncClient(base_url=LLM_BASE_URL, headers=LLM_HEADERS, timeout=http_timeout) as client:
try:
async with client.stream('POST', '/chat/completions', json=payload) as response:
response.raise_for_status()
@@ -507,7 +521,7 @@ async def call_vlm_ocr(image_bytes: bytes, language: str = 'auto') -> str:
http_timeout = httpx.Timeout(connect=10.0, read=None, write=30.0, pool=30.0)
try:
async with httpx.AsyncClient(base_url=LLM_BASE_URL, timeout=http_timeout) as client:
async with httpx.AsyncClient(base_url=LLM_BASE_URL, headers=LLM_HEADERS, timeout=http_timeout) as client:
resp = await asyncio.wait_for(
client.post('/chat/completions', json=payload), timeout=OCR_TIMEOUT,
)