90 lines
2.3 KiB
Python
90 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import ollama
|
|
|
|
|
|
DEFAULT_MODEL = "qwen2.5:3b"
|
|
DEFAULT_FILES = [
|
|
"README.md",
|
|
"Auto_Seat.ino",
|
|
"upload.ps1",
|
|
]
|
|
|
|
|
|
def read_context(paths: list[str]) -> str:
|
|
chunks: list[str] = []
|
|
for item in paths:
|
|
path = Path(item)
|
|
if not path.exists() or not path.is_file():
|
|
continue
|
|
|
|
text = path.read_text(encoding="utf-8", errors="replace")
|
|
chunks.append(f"## {path.as_posix()}\n{text}")
|
|
|
|
return "\n\n".join(chunks)
|
|
|
|
|
|
def ask(model: str, context: str, question: str) -> str:
|
|
system_prompt = (
|
|
"You are a local assistant for this Arduino Auto_Seat project. "
|
|
"Answer in Korean by default. Use the provided project context when it is relevant. "
|
|
"If the context does not contain the answer, say that clearly and reason from general knowledge."
|
|
)
|
|
|
|
if context:
|
|
user_prompt = f"Project context:\n{context}\n\nQuestion:\n{question}"
|
|
else:
|
|
user_prompt = question
|
|
|
|
response = ollama.chat(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
],
|
|
)
|
|
return response["message"]["content"]
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Local LLM chat for the Auto_Seat project.")
|
|
parser.add_argument("--model", default=DEFAULT_MODEL, help=f"Ollama model name. Default: {DEFAULT_MODEL}")
|
|
parser.add_argument(
|
|
"--no-context",
|
|
action="store_true",
|
|
help="Chat without loading project files into the prompt.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
context = "" if args.no_context else read_context(DEFAULT_FILES)
|
|
|
|
print(f"Local LLM ready: {args.model}")
|
|
print("Ask anything. Type /exit to quit.")
|
|
|
|
while True:
|
|
try:
|
|
question = input("\nYou> ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
print()
|
|
break
|
|
|
|
if question.lower() in {"/exit", "/quit", "exit", "quit"}:
|
|
break
|
|
if not question:
|
|
continue
|
|
|
|
try:
|
|
answer = ask(args.model, context, question)
|
|
except Exception as exc:
|
|
print(f"\nError: {exc}")
|
|
continue
|
|
|
|
print(f"\nLLM> {answer}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|