refine code

This commit is contained in:
geekan 2024-01-11 18:08:04 +08:00
parent decde3290b
commit 97ee2f0c98
8 changed files with 112 additions and 91 deletions

View file

@ -94,84 +94,5 @@ class Context(BaseModel):
return llm
class ContextMixin(BaseModel):
"""Mixin class for context and config"""
# Env/Role/Action will use this context as private context, or use self.context as public context
_context: Optional[Context] = None
# Env/Role/Action will use this config as private config, or use self.context.config as public config
_config: Optional[Config] = None
# Env/Role/Action will use this llm as private llm, or use self.context._llm instance
_llm: Optional[BaseLLM] = None
def __init__(
self,
context: Optional[Context] = None,
config: Optional[Config] = None,
llm: Optional[BaseLLM] = None,
**kwargs,
):
"""Initialize with config"""
super().__init__(**kwargs)
self.set_context(context)
self.set_config(config)
self.set_llm(llm)
def set(self, k, v, override=False):
"""Set attribute"""
if override or not self.__dict__.get(k):
self.__dict__[k] = v
def set_context(self, context: Context, override=True):
"""Set context"""
self.set("_context", context, override)
def set_config(self, config: Config, override=False):
"""Set config"""
self.set("_config", config, override)
def set_llm(self, llm: BaseLLM, override=False):
"""Set llm"""
self.set("_llm", llm, override)
@property
def config(self) -> Config:
"""Role config: role config > context config"""
if self._config:
return self._config
return self.context.config
@config.setter
def config(self, config: Config) -> None:
"""Set config"""
self.set_config(config)
@property
def context(self) -> Context:
"""Role context: role context > context"""
if self._context:
return self._context
return CONTEXT
@context.setter
def context(self, context: Context) -> None:
"""Set context"""
self.set_context(context)
@property
def llm(self) -> BaseLLM:
"""Role llm: if not existed, init from role.config"""
# print(f"class:{self.__class__.__name__}({self.name}), llm: {self._llm}, llm_config: {self._llm_config}")
if not self._llm:
self._llm = self.context.llm_with_cost_manager_from_llm_config(self.config.llm)
return self._llm
@llm.setter
def llm(self, llm: BaseLLM) -> None:
"""Set llm"""
self._llm = llm
# Global context, not in Env
CONTEXT = Context()