versionhq 1.2.4.3__tar.gz → 1.2.4.5__tar.gz

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: versionhq
-Version: 1.2.4.3
+Version: 1.2.4.5
 Summary: Autonomous agent networks for task automation with multi-step reasoning.
 Author-email: Kuriko Iwai <kuriko@versi0n.io>
 License: MIT License
@@ -295,7 +295,7 @@ def dummy_func(message: str, test1: str, test2: list[str]) -> str:
 
 task = vhq.Task(
    description="Amazing task",
-   pydantic_output=CustomOutput,
+   response_schema=CustomOutput,
    callback=dummy_func,
    callback_kwargs=dict(message="Hi! Here is the result: ")
 )
@@ -317,13 +317,13 @@ agent_b = vhq.Agent(role="Leader", llm="gemini-2.0")
 
 task_1 = vhq.Task(
    description="Analyze the client's business model.",
-   response_fields=[vhq.ResponseField(title="test1", data_type=str, required=True),],
+   response_schema=[vhq.ResponseField(title="test1", data_type=str, required=True),],
    allow_delegation=True
 )
 
 task_2 = vhq.Task(
    description="Define a cohort.",
-   response_fields=[vhq.ResponseField(title="test1", data_type=int, required=True),],
+   response_schema=[vhq.ResponseField(title="test1", data_type=int, required=True),],
    allow_delegation=False
 )
 

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/README.md

@@ -212,7 +212,7 @@ def dummy_func(message: str, test1: str, test2: list[str]) -> str:
 
 task = vhq.Task(
    description="Amazing task",
-   pydantic_output=CustomOutput,
+   response_schema=CustomOutput,
    callback=dummy_func,
    callback_kwargs=dict(message="Hi! Here is the result: ")
 )
@@ -234,13 +234,13 @@ agent_b = vhq.Agent(role="Leader", llm="gemini-2.0")
 
 task_1 = vhq.Task(
    description="Analyze the client's business model.",
-   response_fields=[vhq.ResponseField(title="test1", data_type=str, required=True),],
+   response_schema=[vhq.ResponseField(title="test1", data_type=str, required=True),],
    allow_delegation=True
 )
 
 task_2 = vhq.Task(
    description="Define a cohort.",
-   response_fields=[vhq.ResponseField(title="test1", data_type=int, required=True),],
+   response_schema=[vhq.ResponseField(title="test1", data_type=int, required=True),],
    allow_delegation=False
 )
 

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/docs/core/llm/index.md

@@ -17,16 +17,17 @@ You can specify a model and integration platform from the list. Else, we'll use
 ```python
 "openai": [
     "gpt-4.5-preview-2025-02-27",
+    "gpt-4",
     "gpt-4o",
     "gpt-4o-mini",
-    "gpt-4",
     "o1-mini",
     "o1-preview",
 ],
 "gemini": [
-    "gemini/gemini-1.5-flash",
-    "gemini/gemini-1.5-pro",
     "gemini/gemini-2.0-flash-exp",
+    "gemini/gemini-2.0-flash",
+    "gemini/gemini-2.0-flash-thinking-exp",
+    "gemini/gemini-2.0-flash-lite-preview-02-05",
 ],
 "anthropic": [
     "claude-3-7-sonnet-latest",
@@ -47,6 +48,7 @@ You can specify a model and integration platform from the list. Else, we'll use
     "bedrock/us.meta.llama3-2-1b-instruct-v1:0",
     "bedrock/us.meta.llama3-2-3b-instruct-v1:0",
     "bedrock/us.meta.llama3-2-11b-instruct-v1:0",
+    "bedrock/us.meta.llama3-2-90b-instruct-v1:0",
     "bedrock/mistral.mistral-7b-instruct-v0:2",
     "bedrock/mistral.mixtral-8x7b-instruct-v0:1",
     "bedrock/mistral.mistral-large-2407-v1:0",

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/docs/core/task/evaluation.md

@@ -59,7 +59,7 @@ class CustomOutput(BaseModel):
 
 task = vhq.Task(
     description="Research a topic to teach a kid aged 6 about math.",
-    pydantic_output=CustomOutput,
+    response_schema=CustomOutput,
     should_evaluate=True, # triggers evaluation
     eval_criteria=["uniquness", "audience fit",],
 

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/docs/core/task/reference.md

@@ -2,30 +2,29 @@
 
 ### Variables
 
-| <div style="width:160px">**Variable**</div> | **Data Type**       | **Default**               | **Description**                   |
-| :---                      | :---                                  | :---                      | :---                              |
-| **`id`**                  | UUID                                  | uuid.uuid4()              | Stores task `id` as an identifier. |
-| **`name`**                | Optional[str]                         | None                      | Stores a task name (Inherited as `node` identifier if the task is dependent) |
-| **`description`**         | str                                   | None                      | Required field to store a concise task description |
-| **`pydantic_output`**     | Optional[Type[BaseModel]]             | None                      | Stores pydantic custom output class for structured response |
-| **`response_fields`**     | Optional[List[ResponseField]]         | list()                    | Stores JSON formats for stuructured response |
-| **`tools`**               | Optional[List[BaseTool \| ToolSet]]   | None                      | Stores tools to be called when the agent executes the task. |
-| **`can_use_agent_tools`** | bool                                  | True                      | Whether to use the agent tools |
-| **`tool_res_as_final`**   | bool                                  | False                     | Whether to make a tool output as a final response from the agent |
-| **`image`**               | Optional[str]                         | None                      | Stores an absolute file path or URL to the image file in string |
-| **`file`**                | Optional[str]                         | None                      | Stores an absolute file path or URL to the file in string |
-| **`audio`**               | Optional[str]                         | None                      | Stores an absolute file path or URL to the audio file in string |
-| **`should_test_run`**     | bool                                  | False                     | Whether to turn on auto-feedback learning             |
-| **`human`**               | bool                                  | False                     | Whether to ask human input during the task execution  |
-| **`execution_type`**      | TaskExecutionType                     | TaskExecutionType.SYNC    | Sync or async execution |
-| **`allow_delegation`**    | bool                                  | False                     | Whether to allow the agent to delegate the task to another agent |
-| **`callback`**            | Optional[Callable]                    | None                      | Callback function to be executed after LLM calling |
-| **`callback_kwargs`**     | Optional[Dict[str, Any]]              | dict()                    | Args for the callback function (if any)|
-| **`should_evaluate`**     | bool                                  | False                     | Whether to evaluate the task output using eval criteria |
-| **`eval_criteria`**       | Optional[List[str]]                   | list()                    | Evaluation criteria given by the human client |
-| **`fsls`**                | Optional[List[str]]                   | None                      | Examples of competitive and/or weak responses |
-| **`processed_agents`**    | Set[str]                              | set()                     | Stores keys of agents that executed the task |
-| **`output`**              | Optional[TaskOutput]                  | None                      | Stores `TaskOutput` object after the execution |
+| <div style="width:160px">**Variable**</div> | **Data Type**       | **Default**       | **Description**                   |
+| :---                      | :---                                  | :---              | :---                              |
+| **`id`**                  | UUID                                  | uuid.uuid4()      | Stores task `id` as an identifier. |
+| **`name`**                | Optional[str]                         | None              | Stores a task name (Inherited as `node` identifier if the task is dependent) |
+| **`description`**         | str                                   | None              | Required field to store a concise task description |
+| **`response_schema`**     | Optional[Type[BaseModel] \| List[ResponseField]] | None   | Response schema for structured output. |
+| **`tools`**               | Optional[List[BaseTool \| ToolSet]]   | None              | Tools, tool sets, or RAG tools |
+| **`can_use_agent_tools`** | bool                                  | True              | Whether to use the agent tools |
+| **`tool_res_as_final`**   | bool                                  | False             | Whether to make a tool output as a final response from the agent |
+| **`image`**               | Optional[str]                         | None              | Absolute file path or URL to the image file |
+| **`file`**                | Optional[str]                         | None              | Absolute file path or URL to the file |
+| **`audio`**               | Optional[str]                         | None              | Absolute file path or URL to the audio file |
+| **`should_test_run`**     | bool                                  | False             | Whether to turn on auto-feedback learning             |
+| **`human`**               | bool                                  | False             | Whether to ask human input during the task execution  |
+| **`execution_type`**      | TaskExecutionType                     | TaskExecutionType.SYNC | Sync or async execution |
+| **`allow_delegation`**    | bool                                  | False             | Whether to allow the agent to delegate the task to another agent |
+| **`callback`**            | Optional[Callable]                    | None              | Callback function to be executed after LLM calling |
+| **`callback_kwargs`**     | Optional[Dict[str, Any]]              | dict()            | Args for the callback function (if any)|
+| **`should_evaluate`**     | bool                                  | False             | Whether to evaluate the task output using eval criteria |
+| **`eval_criteria`**       | Optional[List[str]]                   | list()            | Evaluation criteria given by the human client |
+| **`fsls`**                | Optional[List[str]]                   | None              | Examples of competitive and/or weak responses |
+| **`processed_agents`**    | Set[str]                              | set()             | Stores keys of agents that executed the task |
+| **`output`**              | Optional[TaskOutput]                  | None              | Stores `TaskOutput` object after the execution |
 
 
 ### Class Methods
@@ -68,8 +67,8 @@
 | :---                      | :---                                  | :---          | :---                                                  |
 | **`task_id`**             | UUID                                  | uuid.uuid4()  | Stores task `id` as an identifier.                    |
 | **`raw`**                 | str                                   | None          | Stores response in plane text format. `None` or `""` when the model returned errors.|
-| **`json_dict`**           | Dict[str, Any]                        | None          | Stores response in JSON serializable dictionary. When the system failed formatting or executing tasks without response_fields, `{ output: <res.raw> }` will be returned. |
-| **`pydantic`**            | Type[`BaseModel`]                     | None          | Populates and stores Pydantic class object defined in the `pydantic_output` field if given. |
+| **`json_dict`**           | Dict[str, Any]                        | None          | Stores response in JSON serializable dictionary. |
+| **`pydantic`**            | Type[`BaseModel`]                     | None          | Populates the given response schema in pydantic class |
 | **`tool_output`**         | Optional[Any]                         | None          | Stores results from the tools of the task or agents ONLY when `tool_res_as_final` set as `True`. |
 | **`callback_output`**     | Optional[Any]                         | None          | Stores results from callback functions if any. |
 | **`latency`**             | Optional[float]                       | None          | Stores job latency in milseconds. |

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/docs/core/task/task-output.md

@@ -27,7 +27,7 @@ def summarize_response(message: str, test1: str, test2: list[str]) -> str:
 
 task = vhq.Task(
     description="Research a topic to teach a kid aged 6 about math.",
-    pydantic_output=CustomOutput,
+    response_schema=CustomOutput,
     tools=[dummy_tool],
     callback=summarize_response,
     callback_kwargs=dict(message="Hi! Here is the result: "),

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/docs/core/task/task-strc-response.md

@@ -8,11 +8,11 @@ But you can choose to generate Pydantic class or specifig JSON object as respons
 
 <hr />
 
-## 1. Pydantic
+`[var]`<bold>`response_schema: Optional[Type[BaseModel] | List[ResponseField]] = None`</bold>
 
-`[var]`<bold>`pydantic_output: Optional[Type[BaseModel]] = None`</bold>
+## 1. Pydantic
 
-Add a `custom Pydantic class` as a structured response format to the `pydantic_output` field.
+Add a `custom Pydantic class` to the `response_schema` field to generate a structured response.
 
 The custom class can accept **one layer of a nested child** as you can see in the following code snippet:
 
@@ -44,7 +44,7 @@ class Demo(BaseModel):
 # 2. Define a task
 task = vhq.Task(
     description="generate random output that strictly follows the given format",
-    pydantic_output=Demo,
+    response_schema=Demo,
 )
 
 # 3. Execute
@@ -63,8 +63,6 @@ assert [
 
 ## 2. JSON
 
-`[var]`<bold>`response_fields: List[InstanceOf[ResponseField]] = None`</bold>
-
 Similar to Pydantic, JSON output structure can be defined by using a list of `ResponseField` objects.
 
 The following code snippet demonstrates how to use `ResponseField` to generate output with a maximum of one level of nesting.
@@ -75,16 +73,15 @@ Custom JSON outputs can accept **one layer of a nested child**.
 
 - `demo_response_fields` in the following case is identical to the previous Demo class, except that titles are specified for nested fields.
 
-- Agents generate JSON output by default, whether or not `response_fields` are used.
+- Agents generate JSON output by default, whether or not `response_schema` is given.
 
-- However, response_fields are REQUIRED to specify JSON key titles and data types.
+- However, `response_schema` is required to specify the key value sets.
 
 ```python
 import versionhq as vhq
 
 # 1. Define a list of ResponseField objects.
 demo_response_fields = [
-    # no nesting
     vhq.ResponseField(title="demo_1", data_type=int),
     vhq.ResponseField(title="demo_2", data_type=float),
     vhq.ResponseField(title="demo_3", data_type=str),
@@ -125,16 +122,15 @@ demo_response_fields = [
 # 2. Define a task
 task = vhq.Task(
     description="Output random values strictly following the data type defined in the given response format.",
-    response_fields=demo_response_fields
+    response_schema=demo_response_fields
 )
 
 
 # 3. Execute
 res = task.execute()
 
-assert isinstance(res, vhq.TaskOutput) and res.task_id is task.id
-assert res.raw and res.json and res.pydantic is None
-assert [v and type(v) == task.response_fields[i].data_type for i, (k, v) in enumerate(res.json_dict.items())]
+assert isinstance(res, vhq.TaskOutput)
+assert [v and type(v) == task.response_schema[i].data_type for i, (k, v) in enumerate(res.json_dict.items())]
 ```
 
 * Ref. <a href="/core/task/response-field">`ResponseField`</a> class
@@ -172,7 +168,7 @@ class Sub(BaseModel):
 
 sub_task = vhq.Task(
     description="generate random values that strictly follows the given format.",
-    pydantic_output=Sub
+    response_schema=Sub
 )
 sub_res = sub_task.execute()
 
@@ -190,7 +186,7 @@ def format_response(sub, main1, main2) -> Main:
 # 3. Executes
 main_task = vhq.Task(
     description="generate random values that strictly follows the given format.",
-    pydantic_output=Main,
+    response_schema=Main,
     callback=format_response,
     callback_kwargs=dict(sub=sub_res.json_dict),
 )

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/docs/quickstart.md

@@ -52,7 +52,7 @@ agent = vhq.Agent(role="demo manager")
 
 task = vhq.Task(
    description="Amazing task",
-   pydantic_output=CustomOutput,
+   response_schema=CustomOutput,
    callback=dummy_func,
    callback_kwargs=dict(message="Hi! Here is the result: ")
 )
@@ -88,13 +88,13 @@ agent_b = vhq.Agent(role="agent b", goal="My amazing goals", llm="llm-of-your-ch
 
 task_1 = vhq.Task(
    description="Analyze the client's business model.",
-   response_fields=[vhq.ResponseField(title="test1", data_type=str, required=True),],
+   response_schema=[vhq.ResponseField(title="test1", data_type=str, required=True),],
    allow_delegation=True
 )
 
 task_2 = vhq.Task(
    description="Define a cohort.",
-   response_fields=[vhq.ResponseField(title="test1", data_type=int, required=True),],
+   response_schema=[vhq.ResponseField(title="test1", data_type=int, required=True),],
    allow_delegation=False
 )
 

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/pyproject.toml

@@ -15,7 +15,7 @@ exclude = ["test*", "__pycache__", "*.egg-info"]
 
 [project]
 name = "versionhq"
-version = "1.2.4.3"
+version = "1.2.4.5"
 authors = [{ name = "Kuriko Iwai", email = "kuriko@versi0n.io" }]
 description = "Autonomous agent networks for task automation with multi-step reasoning."
 readme = "README.md"

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/__init__.py

@@ -32,7 +32,7 @@ from versionhq.agent_network.formation import form_agent_network
 from versionhq.task_graph.draft import workflow
 
 
-__version__ = "1.2.4.3"
+__version__ = "1.2.4.5"
 __all__ = [
     "Agent",
 

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/_prompt/auto_feedback.py

@@ -67,7 +67,7 @@ class PromptFeedbackGraph(TaskGraph):
         if not agents:
             return None
 
-        self.concl_template = base_task.pydantic_output if base_task.pydantic_output else base_task.response_fields if base_task.response_fields else None
+        self.concl_response_schema = base_task.response_schema
         base_agent.callbacks.append(self._reflect)
         init_node = Node(task=base_task, assigned_to=base_agent)
         self.add_node(init_node)

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/_prompt/model.py

@@ -2,7 +2,7 @@
 from typing import Dict, List, Tuple, Any
 from textwrap import dedent
 
-from pydantic import InstanceOf
+from pydantic import InstanceOf, BaseModel
 
 from versionhq._utils import is_valid_url
 
@@ -25,34 +25,26 @@ class Prompt:
 
 
     def _draft_output_prompt(self) -> str:
-        """Drafts prompt for output either from `pydantic_output` or `response_fields`"""
+        """Drafts prompt for output format using `response_schema`."""
 
-        from versionhq.llm.model import DEFAULT_MODEL_PROVIDER_NAME
+        from versionhq.task.model import ResponseField
 
         output_prompt = ""
-        model_provider = self.agent.llm.provider if self.agent else DEFAULT_MODEL_PROVIDER_NAME
+        output_formats_to_follow = dict()
 
-        if self.task.pydantic_output:
-            output_prompt, output_formats_to_follow = "", dict()
-            response_format = str(self.task._structure_response_format(model_provider=model_provider))
-            for k, v in self.task.pydantic_output.model_fields.items():
-                output_formats_to_follow[k] = f"<Return your answer in {v.annotation}>"
+        if self.task.response_schema:
+            if isinstance(self.task.response_schema, list):
+                for item in self.task.response_schema:
+                    if isinstance(item, ResponseField):
+                        output_formats_to_follow[item.title] = f"<Return your answer in {item.data_type.__name__}>"
 
-            output_prompt = f"""Your response MUST be a valid JSON string that strictly follows the response format. Use double quotes for all keys and string values. Do not use single quotes, trailing commas, or any other non-standard JSON syntax.
-Response format: {response_format}
-Ref. Output image: {output_formats_to_follow}
-"""
-        elif self.task.response_fields:
-            output_prompt, output_formats_to_follow = "", dict()
-            response_format = str(self.task._structure_response_format(model_provider=model_provider))
-            for item in self.task.response_fields:
-                if item:
-                    output_formats_to_follow[item.title] = f"<Return your answer in {item.data_type.__name__}>"
+            elif issubclass(self.task.response_schema, BaseModel):
+                for k, v in self.task.response_schema.model_fields.items():
+                    output_formats_to_follow[k] = f"<Return your answer in {v.annotation}>"
 
             output_prompt = f"""Your response MUST be a valid JSON string that strictly follows the response format. Use double quotes for all keys and string values. Do not use single quotes, trailing commas, or any other non-standard JSON syntax.
-Response format: {response_format}
 Ref. Output image: {output_formats_to_follow}
-"""
+    """
         else:
             output_prompt = "You MUST return your response as a valid JSON serializable string, enclosed in double quotes. Use double quotes for all keys and string values. Do NOT use single quotes, trailing commas, or other non-standard JSON syntax."
 

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/agent/inhouse_agents.py

@@ -30,7 +30,7 @@ vhq_task_evaluator = Agent(
 vhq_formation_planner = Agent(
     role="vhq-Formation Planner",
     goal="Plan a formation of agents based on the given task descirption.",
-    llm="gemini/gemini-2.0-flash-exp",
+    llm="gemini/gemini-2.0-flash",
     llm_config=dict(top_p=0.8, topK=40, temperature=0.9),
     maxit=1,
     max_retry_limit=1,
@@ -46,7 +46,7 @@ vhq_formation_planner = Agent(
 vhq_agent_creator = Agent(
     role="vhq-Agent Creator",
     goal="build an agent that can handle the given task",
-    llm="gemini/gemini-2.0-flash-exp",
+    llm="gemini/gemini-2.0-flash",
     maxit=1,
     max_retry_limit=1,
 )

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/agent/model.py

@@ -1,6 +1,6 @@
 import os
 import uuid
-from typing import Any, Dict, List, Optional, TypeVar, Callable, Type
+from typing import Any, Dict, List, Optional, TypeVar, Callable, Type, Tuple
 from typing_extensions import Self
 from dotenv import load_dotenv
 
@@ -475,9 +475,16 @@ class Agent(BaseModel):
         return self
 
 
-    def start(self, context: Any = None, tool_res_as_final: bool = False, image: str = None, file: str = None, audio: str = None) -> Any | None:
+    def start(
+            self,
+            context: Any = None,
+            tool_res_as_final: bool = False,
+            image: str = None,
+            file: str = None,
+            audio: str = None
+        ) -> Tuple[Any | None, Any | None]:
         """
-        Defines and executes a task when it is not given and returns TaskOutput object.
+        Defines and executes a task, then returns TaskOutput object with the generated task.
         """
 
         if not self.role:
@@ -491,14 +498,14 @@ class Agent(BaseModel):
 
         task = Task(
             description=f"Generate a simple result in a sentence to achieve the goal: {self.goal if self.goal else self.role}. If needed, list up necessary steps in concise manner.",
-            pydantic_output=Output,
+            response_schema=Output,
             tool_res_as_final=tool_res_as_final,
             image=image, #REFINEME - query memory/knowledge or self create
             file=file,
             audio=audio,
         )
         res = task.execute(agent=self, context=context)
-        return res
+        return res, task
 
 
     def execute_task(self, task, context: Optional[Any] = None, task_tools: Optional[List[Tool | ToolSet]] = list()) -> str:

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/agent_network/formation.py

@@ -74,8 +74,8 @@ def form_agent_network(
         leader_agent: str
 
     vhq_task = Task(
-        description=f"Design a team of specialized agents to fully automate the following task and achieve the expected outcome. For each agent, define its role, task description, and expected outputs via the task with items in a list. Then specify the formation if the formation is not given. If you think SUPERVISING or HYBRID is the best formation, include a leader_agent role, else leave the leader_agent role blank.\nTask: {str(task)}\nExpected outcome: {prompt_expected_outcome}\nFormation: {prompt_formation}",
-        pydantic_output=Outcome
+        description=f"Design a team of specialized agents to fully automate the following task and deliver the expected outcome. For each agent, define its role, task description, and expected outputs via the task with items in a list. Then specify the formation if the formation is not given. If you think SUPERVISING or HYBRID is the best formation, include a leader_agent role, else leave the leader_agent role blank.\nTask: {str(task)}\nExpected outcome: {prompt_expected_outcome}\nFormation: {prompt_formation}",
+        response_schema=Outcome
     )
 
     if agents:
@@ -120,7 +120,7 @@ def form_agent_network(
                     except:
                         pass
                 output = create_model("Output", **fields) if fields else None
-                _task = Task(description=task_descriptions[i], pydantic_output=output)
+                _task = Task(description=task_descriptions[i], response_schema=output)
                 created_tasks.append(_task)
 
     if len(created_tasks) <= len(created_agents):

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/agent_network/model.py

@@ -206,7 +206,7 @@ class AgentNetwork(BaseModel):
         for unassgined_task in unassigned_tasks:
             task = Task(
                 description=f"Based on the following task summary, draft an agent's role and goal in concise manner. Task summary: {unassgined_task.summary}",
-                response_fields=[
+                response_schema=[
                     ResponseField(title="goal", data_type=str, required=True),
                     ResponseField(title="role", data_type=str, required=True),
                 ],

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/llm/llm_vars.py

@@ -27,8 +27,9 @@ MODELS = {
         "o1-preview",
     ],
     "gemini": [
-        "gemini/gemini-1.5-flash",
-        "gemini/gemini-1.5-pro",
+        "gemini/gemini-2.0-flash",
+        "gemini/gemini-2.0-flash-thinking-exp",
+        "gemini/gemini-2.0-flash-lite-preview-02-05",
         "gemini/gemini-2.0-flash-exp",
     ],
     "anthropic": [
@@ -75,6 +76,7 @@ ENV_VARS = {
     "huggingface": ["HUGGINGFACE_API_KEY", ],
     "bedrock": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME"],
     "sagemaker": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME"],
+    "azure_ai": ["AZURE_AI_API_KEY", "AZURE_AI_API_BASE"],
 }
 
 
@@ -90,9 +92,10 @@ LLM_CONTEXT_WINDOW_SIZES = {
     "o1-preview": 128000,
     "o1-mini": 128000,
 
-    "gemini/gemini-1.5-flash": 1048576,
-    "gemini/gemini-1.5-pro": 2097152,
     "gemini/gemini-2.0-flash-exp": 1048576,
+    "gemini/gemini-2.0-flash": 1048576,
+    "gemini/gemini-2.0-flash-thinking-exp": 1048576,
+    "gemini/gemini-2.0-flash-lite-preview-02-05": 1048576,
 
     "claude-3-7-sonnet-latest": 200000,
     "claude-3-5-haiku-latest": 200000,

{versionhq-1.2.4.3 → versionhq-1.2.4.5}/src/versionhq/llm/model.py

@@ -228,6 +228,9 @@ class LLM(BaseModel):
             if self.context_window_size and valid_config["max_tokens"] > self.context_window_size:
                 valid_config["max_tokens"] = self.context_window_size
 
+        if "model" in valid_config:
+            self.model = valid_config.pop("model")
+
         self.llm_config = valid_config
         return valid_config
 
@@ -389,7 +392,6 @@ class LLM(BaseModel):
                     self._usages.append(res["usage"])
                     return res.choices[0].message.content
 
-
             except JSONSchemaValidationError as e:
                 self._logger.log(level="error", message="Raw Response: {}".format(e.raw_response), color="red")
                 raise e