mirror of
https://github.com/tcsenpai/agenticSeek.git
synced 2025-06-03 01:30:11 +00:00
67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
|
|
import threading
|
|
import logging
|
|
from abc import abstractmethod
|
|
from .cache import Cache
|
|
|
|
class GenerationState:
|
|
def __init__(self):
|
|
self.lock = threading.Lock()
|
|
self.last_complete_sentence = ""
|
|
self.current_buffer = ""
|
|
self.is_generating = False
|
|
|
|
def status(self) -> dict:
|
|
return {
|
|
"sentence": self.current_buffer,
|
|
"is_complete": not self.is_generating,
|
|
"last_complete_sentence": self.last_complete_sentence,
|
|
"is_generating": self.is_generating,
|
|
}
|
|
|
|
class GeneratorLLM():
|
|
def __init__(self):
|
|
self.model = None
|
|
self.state = GenerationState()
|
|
self.logger = logging.getLogger(__name__)
|
|
handler = logging.StreamHandler()
|
|
handler.setLevel(logging.INFO)
|
|
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
handler.setFormatter(formatter)
|
|
self.logger.addHandler(handler)
|
|
self.logger.setLevel(logging.INFO)
|
|
cache = Cache()
|
|
|
|
def set_model(self, model: str) -> None:
|
|
self.logger.info(f"Model set to {model}")
|
|
self.model = model
|
|
|
|
def start(self, history: list) -> bool:
|
|
if self.model is None:
|
|
raise Exception("Model not set")
|
|
with self.state.lock:
|
|
if self.state.is_generating:
|
|
return False
|
|
self.state.is_generating = True
|
|
self.logger.info("Starting generation")
|
|
threading.Thread(target=self.generate, args=(history,)).start()
|
|
return True
|
|
|
|
def get_status(self) -> dict:
|
|
with self.state.lock:
|
|
return self.state.status()
|
|
|
|
@abstractmethod
|
|
def generate(self, history: list) -> None:
|
|
"""
|
|
Generate text using the model.
|
|
args:
|
|
history: list of strings
|
|
returns:
|
|
None
|
|
"""
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
generator = GeneratorLLM()
|
|
generator.get_status() |