qa_with_llm.py

import os
import subprocess
import sys
from openai import OpenAI

PROMPT_TEMPLATE = (
    "Based on the above context, write an answer in the form of a short phrase for the following question. "
    "For episodic-type memory, When answering the time an event occurred (or will occur), directly use the timestamp in the brackets preceding the episodic memory. When answering, provide the absolute time whenever possible."
    "Keep the answer as concise as possible."
    "Answer with exact words from the context whenever possible.\n\nQuestion: {query}"
)

def retrieve_context( query):
    """
    调用 main.py 的 retrieve 功能，返回 context_block 文本。
    """
    cmd = [
        sys.executable, "main.py", "retrieve",
        "--query", query
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"Memory retrieval failed: {result.stderr}")
    return result.stdout.strip()

def ask_llm(context, query):
    """
    调用大模型，根据 context 和 query 生成英文短语回答。
    """
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    prompt = context + "\n\n" + PROMPT_TEMPLATE.format(query=query)
    completion = client.chat.completions.create(
        model="qwen-plus",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt},
        ],
    )
    return completion.choices[0].message.content.strip()

def main():
    import argparse
    parser = argparse.ArgumentParser(description="基于大模型的记忆问答")
    parser.add_argument("--query", required=True, help="问题")
    args = parser.parse_args()

    context = retrieve_context(args.query)
    answer = ask_llm(context, args.query)
    print(answer)

if __name__ == "__main__":
    main()