janus-llm 4.0.0__tar.gz → 4.2.0__tar.gz

{janus_llm-4.0.0 → janus_llm-4.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: janus-llm
-Version: 4.0.0
+Version: 4.2.0
 Summary: A transcoding library using LLMs.
 Home-page: https://github.com/janus-llm/janus-llm
 License: Apache 2.0
@@ -23,20 +23,28 @@ Requires-Dist: langchain-anthropic (>=0.1.15,<0.2.0)
 Requires-Dist: langchain-community (>=0.2.0,<0.3.0)
 Requires-Dist: langchain-core (>=0.2.0,<0.3.0)
 Requires-Dist: langchain-openai (>=0.1.8,<0.2.0)
+Requires-Dist: langchain-unstructured (>=0.1.2,<0.2.0)
 Requires-Dist: nltk (>=3.8.1,<4.0.0)
 Requires-Dist: numpy (>=1.24.3,<2.0.0)
 Requires-Dist: openai (>=1.14.0,<2.0.0)
+Requires-Dist: pi-heif (>=0.20.0,<0.21.0)
 Requires-Dist: py-readability-metrics (>=1.4.5,<2.0.0)
 Requires-Dist: py-rouge (>=1.1,<2.0)
+Requires-Dist: pytesseract (>=0.3.13,<0.4.0)
 Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
 Requires-Dist: rich (>=13.7.1,<14.0.0)
 Requires-Dist: sacrebleu (>=2.4.1,<3.0.0)
+Requires-Dist: scikit-learn (>=1.5.2,<2.0.0)
 Requires-Dist: sentence-transformers (>=2.6.1,<3.0.0) ; extra == "hf-local" or extra == "all"
+Requires-Dist: tesseract (>=0.1.3,<0.2.0)
 Requires-Dist: text-generation (>=0.6.0,<0.7.0)
 Requires-Dist: tiktoken (>=0.7.0,<0.8.0)
 Requires-Dist: transformers (>=4.31.0,<5.0.0)
 Requires-Dist: tree-sitter (>=0.21.0,<0.22.0)
 Requires-Dist: typer (>=0.9.0,<0.10.0)
+Requires-Dist: unstructured (>=0.15.9,<0.16.0)
+Requires-Dist: unstructured-inference (>=0.7.36,<0.8.0)
+Requires-Dist: unstructured-pytesseract (>=0.3.13,<0.4.0)
 Project-URL: Documentation, https://janus-llm.github.io/janus-llm
 Project-URL: Repository, https://github.com/janus-llm/janus-llm
 Description-Content-Type: text/markdown

{janus_llm-4.0.0 → janus_llm-4.2.0}/janus/__init__.py

@@ -5,7 +5,7 @@ from langchain_core._api.deprecation import LangChainDeprecationWarning
 from janus.converter.translate import Translator
 from janus.metrics import *  # noqa: F403
 
-__version__ = "4.0.0"
+__version__ = "4.2.0"
 
 # Ignoring a deprecation warning from langchain_core that I can't seem to hunt down
 warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

{janus_llm-4.0.0 → janus_llm-4.2.0}/janus/cli.py

@@ -13,10 +13,13 @@ from rich.console import Console
 from rich.prompt import Confirm
 from typing_extensions import Annotated
 
+import janus.refiners.refiner
+import janus.refiners.uml
 from janus.converter.aggregator import Aggregator
 from janus.converter.converter import Converter
 from janus.converter.diagram import DiagramGenerator
 from janus.converter.document import Documenter, MadLibsDocumenter, MultiDocumenter
+from janus.converter.partition import Partitioner
 from janus.converter.requirements import RequirementsDocumenter
 from janus.converter.translate import Translator
 from janus.embedding.collections import Collections
@@ -39,11 +42,11 @@ from janus.llm.models_info import (
     MODEL_TYPE_CONSTRUCTORS,
     MODEL_TYPES,
     TOKEN_LIMITS,
+    azure_models,
     bedrock_models,
     openai_models,
 )
 from janus.metrics.cli import evaluate
-from janus.refiners.refiner import REFINERS
 from janus.utils.enums import LANGUAGES
 from janus.utils.logger import create_logger
 
@@ -68,6 +71,18 @@ with open(db_file, "r") as f:
 collections_config_file = Path(db_loc) / "collections.json"
 
 
+def get_subclasses(cls):
+    return set(cls.__subclasses__()).union(
+        set(s for c in cls.__subclasses__() for s in get_subclasses(c))
+    )
+
+
+REFINER_TYPES = get_subclasses(janus.refiners.refiner.JanusRefiner).union(
+    {janus.refiners.refiner.JanusRefiner}
+)
+REFINERS = {r.__name__: r for r in REFINER_TYPES}
+
+
 def get_collections_config():
     if collections_config_file.exists():
         with open(collections_config_file, "r") as f:
@@ -243,22 +258,23 @@ def translate(
             click_type=click.Choice(list(CUSTOM_SPLITTERS.keys())),
         ),
     ] = "file",
-    refiner_type: Annotated[
-        str,
+    refiner_types: Annotated[
+        list[str],
         typer.Option(
             "-r",
             "--refiner",
-            help="Name of custom refiner to use",
+            help="List of refiner types to use. Add -r for each refiner to use in\
+                refinement chain",
             click_type=click.Choice(list(REFINERS.keys())),
         ),
-    ] = "none",
+    ] = ["JanusRefiner"],
     retriever_type: Annotated[
         str,
         typer.Option(
             "-R",
             "--retriever",
             help="Name of custom retriever to use",
-            click_type=click.Choice(["active_usings"]),
+            click_type=click.Choice(["active_usings", "language_docs"]),
         ),
     ] = None,
     max_tokens: Annotated[
@@ -271,6 +287,7 @@ def translate(
         ),
     ] = None,
 ):
+    refiner_types = [REFINERS[r] for r in refiner_types]
     try:
         target_language, target_version = target_lang.split("-")
     except ValueError:
@@ -295,7 +312,7 @@ def translate(
         db_path=db_loc,
         db_config=collections_config,
         splitter_type=splitter_type,
-        refiner_type=refiner_type,
+        refiner_types=refiner_types,
         retriever_type=retriever_type,
     )
     translator.translate(input_dir, output_dir, overwrite, collection)
@@ -401,22 +418,23 @@ def document(
             click_type=click.Choice(list(CUSTOM_SPLITTERS.keys())),
         ),
     ] = "file",
-    refiner_type: Annotated[
-        str,
+    refiner_types: Annotated[
+        list[str],
         typer.Option(
             "-r",
             "--refiner",
-            help="Name of custom refiner to use",
+            help="List of refiner types to use. Add -r for each refiner to use in\
+                refinement chain",
             click_type=click.Choice(list(REFINERS.keys())),
         ),
-    ] = "none",
+    ] = ["JanusRefiner"],
     retriever_type: Annotated[
         str,
         typer.Option(
             "-R",
             "--retriever",
             help="Name of custom retriever to use",
-            click_type=click.Choice(["active_usings"]),
+            click_type=click.Choice(["active_usings", "language_docs"]),
         ),
     ] = None,
     max_tokens: Annotated[
@@ -429,6 +447,7 @@ def document(
         ),
     ] = None,
 ):
+    refiner_types = [REFINERS[r] for r in refiner_types]
     model_arguments = dict(temperature=temperature)
     collections_config = get_collections_config()
     kwargs = dict(
@@ -440,7 +459,7 @@ def document(
         db_path=db_loc,
         db_config=collections_config,
         splitter_type=splitter_type,
-        refiner_type=refiner_type,
+        refiner_types=refiner_types,
         retriever_type=retriever_type,
     )
     if doc_mode == "madlibs":
@@ -457,12 +476,6 @@ def document(
     documenter.translate(input_dir, output_dir, overwrite, collection)
 
 
-def get_subclasses(cls):
-    return set(cls.__subclasses__()).union(
-        set(s for c in cls.__subclasses__() for s in get_subclasses(c))
-    )
-
-
 @app.command()
 def aggregate(
     input_dir: Annotated[
@@ -577,6 +590,103 @@ def aggregate(
     aggregator.translate(input_dir, output_dir, overwrite, collection)
 
 
+@app.command(
+    help="Partition input code using an LLM.",
+    no_args_is_help=True,
+)
+def partition(
+    input_dir: Annotated[
+        Path,
+        typer.Option(
+            "--input",
+            "-i",
+            help="The directory containing the source code to be partitioned. ",
+        ),
+    ],
+    language: Annotated[
+        str,
+        typer.Option(
+            "--language",
+            "-l",
+            help="The language of the source code.",
+            click_type=click.Choice(sorted(LANGUAGES)),
+        ),
+    ],
+    output_dir: Annotated[
+        Path,
+        typer.Option(
+            "--output-dir", "-o", help="The directory to store the partitioned code in."
+        ),
+    ],
+    llm_name: Annotated[
+        str,
+        typer.Option(
+            "--llm",
+            "-L",
+            help="The custom name of the model set with 'janus llm add'.",
+        ),
+    ] = "gpt-4o",
+    max_prompts: Annotated[
+        int,
+        typer.Option(
+            "--max-prompts",
+            "-m",
+            help="The maximum number of times to prompt a model on one functional block "
+            "before exiting the application. This is to prevent wasting too much money.",
+        ),
+    ] = 10,
+    overwrite: Annotated[
+        bool,
+        typer.Option(
+            "--overwrite/--preserve",
+            help="Whether to overwrite existing files in the output directory",
+        ),
+    ] = False,
+    temperature: Annotated[
+        float,
+        typer.Option("--temperature", "-t", help="Sampling temperature.", min=0, max=2),
+    ] = 0.7,
+    splitter_type: Annotated[
+        str,
+        typer.Option(
+            "-S",
+            "--splitter",
+            help="Name of custom splitter to use",
+            click_type=click.Choice(list(CUSTOM_SPLITTERS.keys())),
+        ),
+    ] = "file",
+    max_tokens: Annotated[
+        int,
+        typer.Option(
+            "--max-tokens",
+            "-M",
+            help="The maximum number of tokens the model will take in. "
+            "If unspecificed, model's default max will be used.",
+        ),
+    ] = None,
+    partition_token_limit: Annotated[
+        int,
+        typer.Option(
+            "--partition-tokens",
+            "-pt",
+            help="The limit on the number of tokens per partition.",
+        ),
+    ] = 8192,
+):
+    model_arguments = dict(temperature=temperature)
+    kwargs = dict(
+        model=llm_name,
+        model_arguments=model_arguments,
+        source_language=language,
+        max_prompts=max_prompts,
+        max_tokens=max_tokens,
+        splitter_type=splitter_type,
+        partition_token_limit=partition_token_limit,
+    )
+    partitioner = Partitioner(**kwargs)
+    partitioner.translate(input_dir, output_dir, overwrite)
+
+
 @app.command(
     help="Diagram input code using an LLM.",
     no_args_is_help=True,
@@ -666,25 +776,27 @@ def diagram(
             click_type=click.Choice(list(CUSTOM_SPLITTERS.keys())),
         ),
     ] = "file",
-    refiner_type: Annotated[
-        str,
+    refiner_types: Annotated[
+        list[str],
         typer.Option(
             "-r",
             "--refiner",
-            help="Name of custom refiner to use",
+            help="List of refiner types to use. Add -r for each refiner to use in\
+                refinement chain",
             click_type=click.Choice(list(REFINERS.keys())),
         ),
-    ] = "none",
+    ] = ["JanusRefiner"],
     retriever_type: Annotated[
         str,
         typer.Option(
             "-R",
             "--retriever",
             help="Name of custom retriever to use",
-            click_type=click.Choice(["active_usings"]),
+            click_type=click.Choice(["active_usings", "language_docs"]),
         ),
     ] = None,
 ):
+    refiner_types = [REFINERS[r] for r in refiner_types]
     model_arguments = dict(temperature=temperature)
     collections_config = get_collections_config()
     diagram_generator = DiagramGenerator(
@@ -695,7 +807,7 @@ def diagram(
         db_path=db_loc,
         db_config=collections_config,
         splitter_type=splitter_type,
-        refiner_type=refiner_type,
+        refiner_types=refiner_types,
         retriever_type=retriever_type,
         diagram_type=diagram_type,
         add_documentation=add_documentation,
@@ -952,7 +1064,7 @@ def llm_add(
             help="The type of the model",
             click_type=click.Choice(sorted(list(MODEL_TYPE_CONSTRUCTORS.keys()))),
         ),
-    ] = "OpenAI",
+    ] = "Azure",
 ):
     if not MODEL_CONFIG_DIR.exists():
         MODEL_CONFIG_DIR.mkdir(parents=True)
@@ -996,6 +1108,7 @@ def llm_add(
             "model_cost": {"input": in_cost, "output": out_cost},
         }
     elif model_type == "OpenAI":
+        print("DEPRECATED: Use 'Azure' instead. CTRL+C to exit.")
         model_id = typer.prompt(
             "Enter the model ID (list model IDs with `janus llm ls -a`)",
             default="gpt-4o",
@@ -1017,6 +1130,28 @@ def llm_add(
             "token_limit": max_tokens,
             "model_cost": model_cost,
         }
+    elif model_type == "Azure":
+        model_id = typer.prompt(
+            "Enter the model ID (list model IDs with `janus llm ls -a`)",
+            default="gpt-4o",
+            type=click.Choice(azure_models),
+            show_choices=False,
+        )
+        params = dict(
+            # Azure uses the "azure_deployment" key for what we're calling "long_model_id"
+            azure_deployment=MODEL_ID_TO_LONG_ID[model_id],
+            temperature=0.7,
+            n=1,
+        )
+        max_tokens = TOKEN_LIMITS[MODEL_ID_TO_LONG_ID[model_id]]
+        model_cost = COST_PER_1K_TOKENS[MODEL_ID_TO_LONG_ID[model_id]]
+        cfg = {
+            "model_type": model_type,
+            "model_id": model_id,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": model_cost,
+        }
     elif model_type == "BedrockChat" or model_type == "Bedrock":
         model_id = typer.prompt(
             "Enter the model ID (list model IDs with `janus llm ls -a`)",

{janus_llm-4.0.0 → janus_llm-4.2.0}/janus/converter/__init__.py

@@ -2,5 +2,6 @@ from janus.converter.converter import Converter
 from janus.converter.diagram import DiagramGenerator
 from janus.converter.document import Documenter, MadLibsDocumenter, MultiDocumenter
 from janus.converter.evaluate import Evaluator
+from janus.converter.partition import Partitioner
 from janus.converter.requirements import RequirementsDocumenter
 from janus.converter.translate import Translator

{janus_llm-4.0.0 → janus_llm-4.2.0}/janus/converter/_tests/test_translate.py

@@ -90,14 +90,14 @@ class TestDiagramGenerator(unittest.TestCase):
     def setUp(self):
         """Set up the tests."""
         self.diagram_generator = DiagramGenerator(
-            model="gpt-4o",
+            model="gpt-4o-mini",
             source_language="fortran",
             diagram_type="Activity",
         )
 
     def test_init(self):
         """Test __init__ method."""
-        self.assertEqual(self.diagram_generator._model_name, "gpt-4o")
+        self.assertEqual(self.diagram_generator._model_name, "gpt-4o-mini")
         self.assertEqual(self.diagram_generator._source_language, "fortran")
         self.assertEqual(self.diagram_generator._diagram_type, "Activity")
 

{janus_llm-4.0.0 → janus_llm-4.2.0}/janus/converter/converter.py

@@ -6,7 +6,12 @@ from typing import Any
 
 from langchain_core.exceptions import OutputParserException
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import Runnable, RunnableParallel, RunnablePassthrough
+from langchain_core.runnables import (
+    Runnable,
+    RunnableLambda,
+    RunnableParallel,
+    RunnablePassthrough,
+)
 from openai import BadRequestError, RateLimitError
 from pydantic import ValidationError
 
@@ -23,15 +28,14 @@ from janus.language.splitter import (
 from janus.llm.model_callbacks import get_model_callback
 from janus.llm.models_info import MODEL_PROMPT_ENGINES, JanusModel, load_model
 from janus.parsers.parser import GenericParser, JanusParser
-from janus.refiners.refiner import (
-    FixParserExceptions,
-    HallucinationRefiner,
-    JanusRefiner,
-    ReflectionRefiner,
-)
+from janus.refiners.refiner import JanusRefiner
 
 # from janus.refiners.refiner import BasicRefiner, Refiner
-from janus.retrievers.retriever import ActiveUsingsRetriever, JanusRetriever
+from janus.retrievers.retriever import (
+    ActiveUsingsRetriever,
+    JanusRetriever,
+    LanguageDocsRetriever,
+)
 from janus.utils.enums import LANGUAGES
 from janus.utils.logger import create_logger
 
@@ -78,7 +82,7 @@ class Converter:
         protected_node_types: tuple[str, ...] = (),
         prune_node_types: tuple[str, ...] = (),
         splitter_type: str = "file",
-        refiner_type: str | None = None,
+        refiner_types: list[type[JanusRefiner]] = [JanusRefiner],
         retriever_type: str | None = None,
     ) -> None:
         """Initialize a Converter instance.
@@ -105,6 +109,7 @@ class Converter:
                 - None
             retriever_type: The type of retriever to use. Valid values:
                 - "active_usings"
+                - "language_docs"
                 - None
         """
         self._changed_attrs: set = set()
@@ -133,10 +138,11 @@ class Converter:
         self._prompt: ChatPromptTemplate
 
         self._parser: JanusParser = GenericParser()
+        self._base_parser: JanusParser = GenericParser()
         self._combiner: Combiner = Combiner()
 
         self._splitter_type: str
-        self._refiner_type: str | None
+        self._refiner_types: list[type[JanusRefiner]]
         self._retriever_type: str | None
 
         self._splitter: Splitter
@@ -144,7 +150,7 @@ class Converter:
         self._retriever: JanusRetriever
 
         self.set_splitter(splitter_type=splitter_type)
-        self.set_refiner(refiner_type=refiner_type)
+        self.set_refiner_types(refiner_types=refiner_types)
         self.set_retriever(retriever_type=retriever_type)
         self.set_model(model_name=model, **model_arguments)
         self.set_prompt(prompt_template=prompt_template)
@@ -170,7 +176,7 @@ class Converter:
         self._load_model()
         self._load_prompt()
         self._load_retriever()
-        self._load_refiner()
+        self._load_refiner_chain()
         self._load_splitter()
         self._load_vectorizer()
         self._load_chain()
@@ -210,13 +216,13 @@ class Converter:
 
         self._splitter_type = splitter_type
 
-    def set_refiner(self, refiner_type: str | None) -> None:
+    def set_refiner_types(self, refiner_types: list[type[JanusRefiner]]) -> None:
         """Validate and set the refiner type
 
         Arguments:
             refiner_type: the type of refiner to use
         """
-        self._refiner_type = refiner_type
+        self._refiner_types = refiner_types
 
     def set_retriever(self, retriever_type: str | None) -> None:
         """Validate and set the retriever type
@@ -355,48 +361,40 @@ class Converter:
     def _load_retriever(self):
         if self._retriever_type == "active_usings":
             self._retriever = ActiveUsingsRetriever()
+        elif self._retriever_type == "language_docs":
+            self._retriever = LanguageDocsRetriever(self._llm, self._source_language)
         else:
             self._retriever = JanusRetriever()
 
-    @run_if_changed("_refiner_type", "_model_name", "max_prompts", "_parser", "_llm")
-    def _load_refiner(self) -> None:
-        """Load the refiner according to this instance's attributes.
-
-        If the relevant fields have not been changed since the last time this method was
-        called, nothing happens.
-        """
-        if self._refiner_type == "parser":
-            self._refiner = FixParserExceptions(
-                llm=self._llm,
-                parser=self._parser,
-                max_retries=self.max_prompts,
-            )
-        elif self._refiner_type == "reflection":
-            self._refiner = ReflectionRefiner(
-                llm=self._llm,
-                parser=self._parser,
-                max_retries=self.max_prompts,
+    @run_if_changed("_refiner_types", "_model_name", "max_prompts", "_parser")
+    def _load_refiner_chain(self) -> None:
+        self._refiner_chain = RunnableParallel(
+            completion=self._llm,
+            prompt_value=RunnablePassthrough(),
+        )
+        for refiner_type in self._refiner_types[:-1]:
+            # NOTE: Do NOT remove refiner_type=refiner_type from lambda.
+            # Due to lambda capture, must be present or chain will not
+            # be correctly constructed.
+            self._refiner_chain = self._refiner_chain | RunnableParallel(
+                completion=lambda x, refiner_type=refiner_type: refiner_type(
+                    llm=self._llm,
+                    parser=self._base_parser,
+                    max_retries=self.max_prompts,
+                ).parse_completion(**x),
+                prompt_value=lambda x: x["prompt_value"],
             )
-        elif self._refiner_type == "hallucination":
-            self._refiner = HallucinationRefiner(
+        self._refiner_chain = self._refiner_chain | RunnableLambda(
+            lambda x: self._refiner_types[-1](
                 llm=self._llm,
                 parser=self._parser,
                 max_retries=self.max_prompts,
-            )
-        else:
-            self._refiner = JanusRefiner(parser=self._parser)
+            ).parse_completion(**x)
+        )
 
-    @run_if_changed("_parser", "_retriever", "_prompt", "_llm", "_refiner")
+    @run_if_changed("_parser", "_retriever", "_prompt", "_llm", "_refiner_chain")
     def _load_chain(self):
-        self.chain = (
-            self._input_runnable()
-            | self._prompt
-            | RunnableParallel(
-                completion=self._llm,
-                prompt_value=RunnablePassthrough(),
-            )
-            | self._refiner.parse_runnable
-        )
+        self.chain = self._input_runnable() | self._prompt | self._refiner_chain
 
     def _input_runnable(self) -> Runnable:
         return RunnableParallel(

janus_llm-4.2.0/janus/converter/partition.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+
+from janus.converter.converter import Converter
+from janus.language.block import TranslatedCodeBlock
+from janus.parsers.partition_parser import PartitionParser
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+class Partitioner(Converter):
+    def __init__(self, partition_token_limit: int, **kwargs):
+        super().__init__(**kwargs)
+        self.set_prompt("partition")
+        self._load_model()
+        self._parser = PartitionParser(
+            token_limit=partition_token_limit,
+            model=self._llm,
+        )
+        self._target_language = self._source_language
+        self._target_suffix = self._source_suffix
+        self._load_parameters()
+
+    def _save_to_file(self, block: TranslatedCodeBlock, out_path: Path) -> None:
+        output_str = self._parser.parse_combined_output(block.complete_text)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        out_path.write_text(output_str, encoding="utf-8")

{janus_llm-4.0.0 → janus_llm-4.2.0}/janus/language/combine.py

@@ -1,3 +1,5 @@
+import re
+
 from janus.language.block import CodeBlock, TranslatedCodeBlock
 from janus.language.file import FileManager
 from janus.utils.logger import create_logger
@@ -90,3 +92,23 @@ class ChunkCombiner(Combiner):
             root: The functional code block to combine with its children.
         """
         return root
+
+
+class PartitionCombiner(Combiner):
+    @staticmethod
+    def combine(root: CodeBlock) -> None:
+        """A combiner which inserts partition tags between code blocks"""
+        queue = [root]
+        while queue:
+            block = queue.pop(0)
+            if block.children:
+                queue.extend(block.children)
+            else:
+                block.affixes = (block.prefix, block.suffix + "\n<JANUS_PARTITION>\n")
+
+        super(PartitionCombiner, PartitionCombiner).combine(root)
+        root.text = re.sub(r"(?:\n<JANUS_PARTITION>\n)+$", "", root.text)
+        root.affixes = (
+            root.prefix,
+            re.sub(r"(?:\n<JANUS_PARTITION>\n)+$", "", root.suffix),
+        )

{janus_llm-4.0.0 → janus_llm-4.2.0}/janus/llm/model_callbacks.py

@@ -23,6 +23,11 @@ openai_model_reroutes = {
     "gpt-3.5-turbo-16k-0613": "gpt-3.5-turbo-0125",
 }
 
+azure_model_reroutes = {
+    "gpt-4o": "gpt-4o-2024-08-06",
+    "gpt-4o-mini": "gpt-4o-mini",
+    "gpt-3.5-turbo-16k": "gpt35-turbo-16k",
+}
 
 # Updated 2024-06-21
 COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
@@ -31,6 +36,10 @@ COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
     "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
     "gpt-4-0613": {"input": 0.03, "output": 0.06},
     "gpt-4o-2024-05-13": {"input": 0.005, "output": 0.015},
+    "gpt-4o-2024-08-06": {"input": 0.00275, "output": 0.011},
+    "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
+    "gpt35-turbo-16k": {"input": 0.003, "output": 0.004},
+    "gpt-35-turbo-16k": {"input": 0.003, "output": 0.004},
     "anthropic.claude-v2": {"input": 0.008, "output": 0.024},
     "anthropic.claude-instant-v1": {"input": 0.0008, "output": 0.0024},
     "anthropic.claude-3-haiku-20240307-v1:0": {"input": 0.00025, "output": 0.00125},