Adolfo Reyna
76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
from .base import LLMEngine
|
|
import subprocess, os, logging, re, tempfile
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
class GeminiCLILLM(LLMEngine):
|
|
def __init__(self, yolo=True, workspace_dir="data", system_file="data/soul.md"):
|
|
self.yolo = yolo
|
|
self.workspace_dir = workspace_dir
|
|
self.system_file = system_file
|
|
self.session_id = None
|
|
|
|
def _read_system(self):
|
|
try:
|
|
if os.path.exists(self.system_file):
|
|
with open(self.system_file) as f:
|
|
return f.read()
|
|
except Exception as e:
|
|
log.warning(f"No system file: {e}")
|
|
return ""
|
|
|
|
def chat(self, messages, speaker=None):
|
|
system = ""
|
|
user_prompt = ""
|
|
for m in messages:
|
|
if m["role"] == "system":
|
|
system += m["content"] + "\n"
|
|
elif m["role"] == "user":
|
|
user_prompt += f"User: {m['content']}\n"
|
|
elif m["role"] == "assistant":
|
|
user_prompt += f"Assistant: {m['content']}\n"
|
|
|
|
if speaker and speaker.get("name") != "unknown":
|
|
user_prompt = f"[Speaker: {speaker.get('name')} conf {speaker.get('fused_conf',0):.2f}]\n" + user_prompt
|
|
|
|
soul = self._read_system()
|
|
full_system = (soul + "\n" + system).strip() if soul else system
|
|
|
|
env = os.environ.copy()
|
|
tmp_sys = None
|
|
if full_system:
|
|
fd, tmp_path = tempfile.mkstemp(suffix=".md")
|
|
os.write(fd, full_system.encode())
|
|
os.close(fd)
|
|
tmp_sys = tmp_path
|
|
env["GEMINI_SYSTEM_MD"] = tmp_sys
|
|
|
|
args = ["gemini", "--prompt", user_prompt]
|
|
if self.yolo:
|
|
args.append("--yolo")
|
|
if self.session_id:
|
|
args.extend(["--resume", self.session_id])
|
|
|
|
try:
|
|
os.makedirs(self.workspace_dir, exist_ok=True)
|
|
proc = subprocess.run(args, capture_output=True, text=True, cwd=self.workspace_dir, env=env, timeout=120)
|
|
out = proc.stdout.strip() or proc.stderr.strip()
|
|
if not self.session_id:
|
|
try:
|
|
res = subprocess.run(["gemini","--list-sessions"], capture_output=True, text=True, cwd=self.workspace_dir, timeout=10)
|
|
matches = re.findall(r"\[([a-f0-9\-]+)\]", res.stdout)
|
|
if matches:
|
|
self.session_id = matches[-1]
|
|
except:
|
|
pass
|
|
return out or "No response from Gemini"
|
|
except Exception as e:
|
|
log.error(f"Gemini CLI error: {e}")
|
|
return f"Gemini error: {e}"
|
|
finally:
|
|
if tmp_sys and os.path.exists(tmp_sys):
|
|
try:
|
|
os.remove(tmp_sys)
|
|
except:
|
|
pass
|