71 lines
2.9 KiB
Python
71 lines
2.9 KiB
Python
|
# utils/llm.py
|
||
|
|
|
||
|
|
import aiohttp
|
||
|
|
import datetime
|
||
|
|
import globals
|
||
|
|
|
||
|
|
from utils.core import miku_vectorstore, miku_lyrics_vectorstore
|
||
|
|
from collections import defaultdict
|
||
|
|
|
||
|
|
conversation_history = defaultdict(list)
|
||
|
|
|
||
|
|
async def query_ollama(user_prompt, user_id):
|
||
|
|
relevant_docs_lore = miku_vectorstore.similarity_search(user_prompt, k=3)
|
||
|
|
relevant_docs_lyrics = miku_lyrics_vectorstore.similarity_search(user_prompt, k=3)
|
||
|
|
|
||
|
|
context_lore = "\n".join([doc.page_content for doc in relevant_docs_lore])
|
||
|
|
context_lyrics = "\n".join([doc.page_content for doc in relevant_docs_lyrics])
|
||
|
|
|
||
|
|
combined_docs = relevant_docs_lore + relevant_docs_lyrics
|
||
|
|
context = "\n\n".join([doc.page_content for doc in combined_docs])
|
||
|
|
|
||
|
|
# Persona definition
|
||
|
|
with open("miku_prompt.txt", "r", encoding="utf-8") as f:
|
||
|
|
system_prompt = f.read()
|
||
|
|
|
||
|
|
# Append angry wake-up note if JUST_WOKEN_UP flag is set
|
||
|
|
if getattr(globals, 'JUST_WOKEN_UP', False):
|
||
|
|
system_prompt += (
|
||
|
|
"\n\n🚨IMPORTANT NOTE🚨: Miku was just rudely woken up by the user with the current prompt of theirs.\nShe is fuming with anger and should respond harshly and coldly.\nDo NOT show kindness or friendliness in this response. Feel free to put certain words in all capital letters to emphasize your frustration."
|
||
|
|
)
|
||
|
|
|
||
|
|
# Add angry wake-up context if in forced angry state
|
||
|
|
if globals.FORCED_ANGRY_UNTIL:
|
||
|
|
now = datetime.datetime.utcnow()
|
||
|
|
if now < globals.FORCED_ANGRY_UNTIL:
|
||
|
|
system_prompt += (
|
||
|
|
"\n\n[NOTE]: Miku is currently angry because she was rudely woken up from sleep by the user. "
|
||
|
|
"Her responses should reflect irritation and coldness towards the user."
|
||
|
|
)
|
||
|
|
|
||
|
|
# Build conversation history
|
||
|
|
history = list(globals.conversation_history.get(user_id, []))[-8:] # limit to last 8 exchanges
|
||
|
|
history_text = "\n".join([f"User: {u}\nMiku: {m}" for u, m in history])
|
||
|
|
|
||
|
|
# Combine prompt
|
||
|
|
full_prompt = (
|
||
|
|
f"{context_lore}\n\n{context_lyrics}\n\n"
|
||
|
|
f"{history_text}\nMiku is currently feeling: {globals.CURRENT_MOOD}\nPlease respond in a way that reflects this emotional tone.\nUser: {user_prompt}\nMiku:"
|
||
|
|
)
|
||
|
|
|
||
|
|
globals.LAST_FULL_PROMPT = full_prompt # ← track latest prompt
|
||
|
|
|
||
|
|
headers = {'Content-Type': 'application/json'}
|
||
|
|
payload = {
|
||
|
|
"model": globals.OLLAMA_MODEL,
|
||
|
|
"prompt": full_prompt,
|
||
|
|
"system": system_prompt,
|
||
|
|
"stream": False
|
||
|
|
}
|
||
|
|
|
||
|
|
async with aiohttp.ClientSession() as session:
|
||
|
|
async with session.post(f"{globals.OLLAMA_URL}/api/generate", json=payload, headers=headers) as response:
|
||
|
|
if response.status == 200:
|
||
|
|
data = await response.json()
|
||
|
|
reply = data.get("response", "No response.")
|
||
|
|
# Save to conversation history
|
||
|
|
globals.conversation_history[user_id].append((user_prompt, reply))
|
||
|
|
return reply
|
||
|
|
else:
|
||
|
|
return f"Error: {response.status}"
|