MOAI / llm_ollama.py
code-slicer's picture
Create llm_ollama.py
c8b7d5d verified
# llm_ollama.py νŒŒμΌμ— λ‹€μŒ λ‚΄μš©μ„ μΆ”κ°€ν•©λ‹ˆλ‹€.
import json
import requests
import streamlit as st
OLLAMA_HOST = "http://localhost:11434"
OLLAMA_MODEL = "gemma2:9b"
OLLAMA_TIMEOUT = 60
KOREAN_SYSTEM_PROMPT = """당신은 ν•œκ΅­μ–΄ μ–΄μ‹œμŠ€ν„΄νŠΈμž…λ‹ˆλ‹€. 항상 ν•œκ΅­μ–΄λ‘œ λ‹΅ν•˜μ„Έμš”."""
STRUCTURED_EXTRACTION_SYSTEM = """\
You are a travel assistant that extracts structured fields from Korean user queries.
Return ONLY a valid JSON object:
{
"emotion": "happy|sad|stressed|excited|tired|none",
"intent": "beach|hiking|shopping|food|museum|relaxing|none",
"country_hint": "",
"city_hint": "",
"themes_hint": ["<0..3 words>"],
"notes": "<very short reasoning in Korean>"
}
If unknown, use "none" or "" and NEVER add extra text outside JSON.
"""
def _call_ollama_chat(messages, model=OLLAMA_MODEL,
temperature=0.8, top_p=0.9, top_k=40, repeat_penalty=1.1,
system_prompt=None):
url = f"{OLLAMA_HOST}/api/chat"
_msgs = []
if system_prompt:
_msgs.append({"role":"system","content":system_prompt})
_msgs.extend(messages)
payload = {
"model": model,
"messages": _msgs,
"options": {
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"repeat_penalty": repeat_penalty,
},
"stream": False,
}
try:
r = requests.post(url, json=payload, timeout=OLLAMA_TIMEOUT)
r.raise_for_status()
return (r.json().get("message") or {}).get("content","") or ""
except requests.RequestException:
return ""
def llm_followup_loop():
# Streamlit은 맀번 μƒˆλ‘œ μ‹€ν–‰λ˜λ―€λ‘œ, λŒ€ν™” 기둝을 μ„Έμ…˜ μƒνƒœμ— μ €μž₯ν•΄μ•Ό 함
if "llm_history" not in st.session_state:
st.session_state.llm_history = [{"role": "bot", "content": "μ–΄λ–€ 점이 더 κΆκΈˆν•˜μ‹ κ°€μš”?"}]
# μ±— λ©”μ‹œμ§€ 기둝
for message in st.session_state.llm_history:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# μ‚¬μš©μž μž…λ ₯ λ°›κΈ°
q = st.chat_input("여행에 λŒ€ν•΄ 더 κΆκΈˆν•œ 점이 μžˆλ‚˜μš”?")
if q:
# μ‚¬μš©μž λ©”μ‹œμ§€ ν‘œμ‹œ 및 기둝
st.session_state.llm_history.append({"role": "user", "content": q})
with st.chat_message("user"):
st.markdown(q)
# Ollama API 호좜
with st.spinner("생각 쀑..."):
messages = [{"role": m["role"], "content": m["content"]} for m in st.session_state.llm_history]
a = _call_ollama_chat(
messages,
system_prompt=KOREAN_SYSTEM_PROMPT,
temperature=0.8, top_p=0.9, top_k=40, repeat_penalty=1.1
)
# 봇 λ©”μ‹œμ§€ ν‘œμ‹œ 및 기둝
if not a:
st.session_state.llm_history.append({"role": "assistant", "content": "⚠️ Ollama 응닡을 λ°›μ§€ λͺ»ν–ˆμŠ΅λ‹ˆλ‹€. μ„œλ²„ μƒνƒœλ₯Ό ν™•μΈν•˜μ„Έμš”."})
else:
st.session_state.llm_history.append({"role": "assistant", "content": a})
# νŽ˜μ΄μ§€ λ¦¬λ‘œλ“œ
st.rerun()