Function run

optillm/plugins/memory_plugin.py:93–134

(system_prompt: str, initial_query: str, client, model: str)

    return margins, response.usage.completion_tokens

def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str, int]:
    memory = Memory()
    query, context = extract_query(initial_query)
    completion_tokens = 0

    # Process context and add to memory
    chunk_size = 100000
    for i in range(0, len(context), chunk_size):
        chunk = context[i:i+chunk_size]
        # print(f"chunk: {chunk}")
        key_info, tokens = extract_key_information(system_prompt, chunk, query, client, model)
        #print(f"key info: {key_info}")
        completion_tokens += tokens
        for info in key_info:
            memory.add(info)
    # print(f"query : {query}")
    # Retrieve relevant information from memory
    relevant_info = memory.get_relevant(query)
    # print(f"relevant_info : {relevant_info}")
    # Generate response using relevant information
    messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"""

I asked my assistant to read and analyse the above content page by page to help you complete this task. These are margin notes left on each page:
'''text
{relevant_info}
'''
Read again the note(s), take a deep breath and answer the query.
{query}
"""}
    ]

    response = client.chat.completions.create(
        model=model,
        messages=messages,
    )
    # print(f"response : {response}")
    final_response = response.choices[0].message.content.strip()
    completion_tokens += response.usage.completion_tokens
    # print(f"final_response: {final_response}")
    return final_response, completion_tokens