backend.law_model module

class backend.law_model.HuggingFaceLLMModified(context_window: int = 3900, max_new_tokens: int = 256, query_wrapper_prompt: str | PromptTemplate = '{query_str}', tokenizer_name: str = 'StabilityAI/stablelm-tuned-alpha-3b', model_name: str = 'StabilityAI/stablelm-tuned-alpha-3b', model: Any | None = None, tokenizer: Any | None = None, device_map: str | None = 'auto', stopping_ids: List[int] | None = None, tokenizer_kwargs: dict | None = None, tokenizer_outputs_to_remove: list | None = None, model_kwargs: dict | None = None, generate_kwargs: dict | None = None, is_chat_model: bool | None = False, callback_manager: CallbackManager | None = None, system_prompt: str = '', messages_to_prompt: Callable[[Sequence[ChatMessage]], str] | None = None, completion_to_prompt: Callable[[str], str] | None = None, pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT, output_parser: BaseOutputParser | None = None)

Bases: HuggingFaceLLM

A modified HuggingFace LLM for complete

complete(prompt: str, formatted: bool = False, **kwargs: Any) → CompletionResponse

Completion endpoint.

class backend.law_model.SimpleLLMAgentWorker(tools: Sequence[BaseTool], llm: LLM, callback_manager: CallbackManager | None = None, verbose: bool = False, tool_retriever: ObjectRetriever[BaseTool] | None = None, *, prefix_message: List[ChatMessage] = None)

Bases: CustomSimpleAgentWorker

A simple agent worker that uses a LLM for chat

finalize_task(task: Task, **kwargs: Any) → None

Finalize task, after all the steps are completed.

get_all_messages(task: Task) → List[ChatMessage]

initialize_step(task: Task, **kwargs: Any) → TaskStep

Initialize step from task.

prefix_message: List[ChatMessage]

backend.law_model.prepare_law_agent(model_name: str, model_url: str | None = None, verbose: bool = False, **kwargs)

Prepare a law agent with the given model name

Params model_name:

the model name

Params model_url:

the model url; if passed, use TGI, else load locally

Params verbose:

verbose mode

Params kwargs:

additional arguments for as_agent

Returns:

a law agent

backend.law_model.prepare_law_llm(model_name: str = 'Equall/Saul-7B-Instruct-v1') → HuggingFaceLLMModified

Prepare a law LLM with the given model name

Params model_name:

the model name

Returns:

a law LLM