Spaces:
Running
Running
File size: 2,629 Bytes
4796bbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | # memory.py
"""
Statelesss-server-friendly conversation memory.
History is passed from the client on each request (no server-side session state)
Compression kicks in when history exceeds the token budget
"""
import logging
import tiktoken
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from .config import get_settings
from .prompt import STANDALONE_QUESTION_PROMPT
logger = logging.getLogger(__name__)
settings = get_settings()
_enc = tiktoken.encoding_for_model("gpt-4o")
def count_tokens(text: str) -> int:
return len(_enc.encode(text))
def build_lc_messages(
history: list[dict],
system_prompt: str,
) -> list:
"""Convert raw history dicts to Langchain message objects"""
messages = [SystemMessage(content=system_prompt)]
for turn in history:
if turn["role"] == "user":
messages.append(HumanMessage(content=turn["content"]))
else:
messages.append(AIMessage(content=turn["content"]))
return messages
def trim_history_to_budget(
history: list[dict],
max_tokens: int = None,
) -> list[dict]:
"""
Sliding window: keep the MOST recent turns that fit the token budget.
Always keeps at minimum the last 2 turns (1 exchange)
"""
budget = max_tokens or settings.context_window_tokens // 3 # 1/3 of budget for history
trimmed: list[dict] = []
total = 0
for turn in reversed(history[-settings.max_history_turns * 2:]): # multiplied by 2 to take both user and AIResponse
tokens = count_tokens(turn["content"])
if total + tokens > budget and len(trimmed) >= 2: # checks if token exceeds budget limit or trimmed is more than 2, to avoid only user or ai going with sys prompt
break
trimmed.insert(0,turn)
total += tokens
return trimmed
async def resolve_standalone_question(
question: str,
history: list[dict],
llm: ChatOpenAI
) -> str:
"""
If conversation history exists, use LLM to rewrite the followup question
as a self-contained query (critical for multi-turn retrieval accuracy)
"""
if not history:
return question
history_str = "\n".join(
f"{t['role'].capitalize()}: {t['content']}" for t in history[-6:]
)
prompt = STANDALONE_QUESTION_PROMPT.format(
history=history_str,
question=question
)
response = await llm.ainvoke([HumanMessage(content=prompt)])
standalone = response.content.strip()
logger.debug(f"Standalone question: '{standalone}")
return standalone
print("[memory] Module ready") |