langroid 0.1.52__py3-none-any.whl → 0.1.54__py3-none-any.whl

langroid/agent/base.py

@@ -1,9 +1,20 @@
 import inspect
 import json
 import logging
+import textwrap
 from abc import ABC
 from contextlib import ExitStack
-from typing import Callable, Dict, List, Optional, Set, Tuple, Type, cast, no_type_check
+from typing import (
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    cast,
+    no_type_check,
+)
 
 from pydantic import BaseSettings, ValidationError
 from rich import print
@@ -15,6 +26,9 @@ from langroid.agent.tool_message import INSTRUCTION, ToolMessage
 from langroid.language_models.base import (
     LanguageModel,
     LLMConfig,
+    LLMMessage,
+    LLMResponse,
+    LLMTokenUsage,
 )
 from langroid.mytypes import DocMetaData, Entity
 from langroid.parsing.json import extract_top_level_json
@@ -60,6 +74,8 @@ class Agent(ABC):
         self.llm_tools_map: Dict[str, Type[ToolMessage]] = {}
         self.llm_tools_handled: Set[str] = set()
         self.llm_tools_usable: Set[str] = set()
+        self.total_llm_token_cost = 0.0
+        self.total_llm_token_usage = 0
         self.default_human_response: Optional[str] = None
         self._indent = ""
         self.llm = LanguageModel.create(config.llm)
@@ -315,7 +331,7 @@ class Agent(ABC):
         else:
             user_msg = Prompt.ask(
                 f"[blue]{self.indent}Human "
-                f"(respond or q, x to exit current level, "
+                "(respond or q, x to exit current level, "
                 f"or hit enter to continue)\n{self.indent}",
             ).strip()
 
@@ -410,6 +426,7 @@ class Agent(ABC):
             if self.llm.get_stream():
                 console.print(f"[green]{self.indent}", end="")
             response = self.llm.generate(prompt, output_len)
+
         displayed = False
         if not self.llm.get_stream() or response.cached:
             # we would have already displayed the msg "live" ONLY if
@@ -417,7 +434,7 @@ class Agent(ABC):
             console.print(f"[green]{self.indent}", end="")
             print("[green]" + response.message)
             displayed = True
-
+        self.update_token_usage(response, prompt, self.llm.get_stream())
         return ChatDocument.from_LLMResponse(response, displayed)
 
     def get_tool_messages(self, msg: str | ChatDocument) -> List[ToolMessage]:
@@ -594,10 +611,66 @@ class Agent(ABC):
             result = f"Error in tool/function-call {tool_name} usage: {type(e)}: {e}"
         return result  # type: ignore
 
-    def num_tokens(self, prompt: str) -> int:
+    def num_tokens(self, prompt: str | List[LLMMessage]) -> int:
         if self.parser is None:
             raise ValueError("Parser must be set, to count tokens")
-        return self.parser.num_tokens(prompt)
+        if isinstance(prompt, str):
+            return self.parser.num_tokens(prompt)
+        else:
+            return sum([self.parser.num_tokens(m.content) for m in prompt])
+
+    def update_token_usage(
+        self, response: LLMResponse, prompt: str | List[LLMMessage], stream: bool
+    ) -> None:
+        """
+        Updates `response.usage` obj (token usage and cost fields).the usage memebr
+        It updates the cost after checking the cache and updates the
+        tokens (prompts and completion) if the response stream is True, because OpenAI
+        doesn't returns these fields.
+
+        Args:
+            response (LLMResponse): LLMResponse object
+            prompt (str | List[LLMMessage]): prompt or list of LLMMessage objects
+            stream (bool): whether to update the usage in the response object
+                if the response is not cached.
+        """
+        if response is not None:
+            # Note: If response was not streamed, then
+            # `response.usage` would already have been set by the API,
+            # so we only need to update in the stream case.
+            if stream:
+                # usage, cost = 0 when response is from cache
+                prompt_tokens = 0
+                completion_tokens = 0
+                cost = 0.0
+                if not response.cached:
+                    prompt_tokens = self.num_tokens(prompt)
+                    completion_tokens = self.num_tokens(response.message)
+                    cost = self.compute_token_cost(prompt_tokens, completion_tokens)
+                response.usage = LLMTokenUsage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    cost=cost,
+                )
+
+            if settings.debug and response.usage is not None:
+                print(
+                    textwrap.dedent(
+                        f"""
+                        Stream: {stream}
+                        prompt_tokens: {response.usage.prompt_tokens}
+                        completion_tokens: {response.usage.completion_tokens}
+                        """.lstrip()
+                    )
+                )
+            # update total counters
+            if response.usage is not None:
+                self.total_llm_token_cost += response.usage.cost
+                self.total_llm_token_usage += response.usage.total_tokens
+
+    def compute_token_cost(self, prompt: int, completion: int) -> float:
+        price = cast(LanguageModel, self.llm).chat_cost()
+        return (price[0] * prompt + price[1] * completion) / 1000
 
     def ask_agent(
         self,

langroid/agent/chat_agent.py

@@ -453,7 +453,8 @@ class ChatAgent(Agent):
             else:
                 response_str = response.message
             print(cached + "[green]" + response_str)
-
+        stream = self.llm.get_stream()  # type: ignore
+        self.update_token_usage(response, messages, stream)
         return ChatDocument.from_LLMResponse(response, displayed)
 
     def _llm_response_temp_context(self, message: str, prompt: str) -> ChatDocument:

langroid/agent/chat_document.py

@@ -7,6 +7,7 @@ from langroid.language_models.base import (
     LLMFunctionCall,
     LLMMessage,
     LLMResponse,
+    LLMTokenUsage,
     Role,
 )
 from langroid.mytypes import DocMetaData, Document, Entity
@@ -29,7 +30,7 @@ class ChatDocMetaData(DocMetaData):
     block: None | Entity = None
     sender_name: str = ""
     recipient: str = ""
-    usage: int = 0
+    usage: Optional[LLMTokenUsage]
     cached: bool = False
     displayed: bool = False
 
@@ -119,7 +120,8 @@ class ChatDocument(Document):
 
     @staticmethod
     def from_LLMResponse(
-        response: LLMResponse, displayed: bool = False
+        response: LLMResponse,
+        displayed: bool = False,
     ) -> "ChatDocument":
         recipient, message = response.get_recipient_and_message()
         return ChatDocument(
@@ -183,7 +185,10 @@ class ChatDocument(Document):
             content = message
 
         return LLMMessage(
-            role=sender_role, content=content, function_call=fun_call, name=sender_name
+            role=sender_role,
+            content=content,
+            function_call=fun_call,
+            name=sender_name,
         )
 
 

langroid/agent/special/doc_chat_agent.py

@@ -7,7 +7,7 @@ Functionality includes:
 """
 import logging
 from contextlib import ExitStack
-from typing import List, Optional, no_type_check
+from typing import List, Optional, Tuple, no_type_check
 
 from rich import print
 from rich.console import Console
@@ -304,7 +304,7 @@ class DocChatAgent(ChatAgent):
         )
 
     @no_type_check
-    def get_relevant_extracts(self, query: str) -> List[Document]:
+    def get_relevant_extracts(self, query: str) -> Tuple[str, List[Document]]:
         """
         Get list of docs or extracts relevant to a query. These could be:
         - the original docs, if they exist and are not too long, or
@@ -316,6 +316,7 @@ class DocChatAgent(ChatAgent):
             query (str): query to search for
 
         Returns:
+            query (str): stand-alone version of input query
             List[Document]: list of relevant docs
 
         """
@@ -341,20 +342,18 @@ class DocChatAgent(ChatAgent):
                     k=self.config.parsing.n_similar_docs,
                 )
             if len(docs_and_scores) == 0:
-                return []
+                return query, []
             passages = [
                 Document(content=d.content, metadata=d.metadata)
                 for (d, _) in docs_and_scores
             ]
 
-        # if passages not too long, no need to extract relevant verbatim text
-        extracts = passages
-        if self.doc_length(passages) > self.config.max_context_tokens:
-            with console.status("[cyan]LLM Extracting verbatim passages..."):
-                with StreamingIfAllowed(self.llm, False):
-                    extracts = self.llm.get_verbatim_extracts(query, passages)
+        with console.status("[cyan]LLM Extracting verbatim passages..."):
+            with StreamingIfAllowed(self.llm, False):
+                extracts = self.llm.get_verbatim_extracts(query, passages)
+                extracts = [e for e in extracts if e.content != NO_ANSWER]
 
-        return extracts
+        return query, extracts
 
     @no_type_check
     def answer_from_docs(self, query: str) -> Document:
@@ -373,7 +372,8 @@ class DocChatAgent(ChatAgent):
                 source="None",
             ),
         )
-        extracts = self.get_relevant_extracts(query)
+        # query may be updated to a stand-alone version
+        query, extracts = self.get_relevant_extracts(query)
         if len(extracts) == 0:
             return response
         with ExitStack() as stack:

langroid/agent/task.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Callable, Dict, List, Optional, Type, cast
+from typing import Callable, Dict, List, Optional, Set, Type, cast
 
 from rich import print
 
@@ -155,7 +155,8 @@ class Task:
 
         # other sub_tasks this task can delegate to
         self.sub_tasks: List[Task] = []
-        self.parent_task: Optional[Task] = None
+        self.parent_task: Set[Task] = set()
+        self.caller: Task | None = None  # which task called this task's `run` method
 
     def __repr__(self) -> str:
         return f"{self.name}"
@@ -165,10 +166,9 @@ class Task:
 
     @property
     def _level(self) -> int:
-        if self.parent_task is None:
+        if self.caller is None:
             return 0
-        else:
-            return self.parent_task._level + 1
+        return self.caller._level + 1
 
     @property
     def _indent(self) -> str:
@@ -199,7 +199,7 @@ class Task:
             return
         assert isinstance(task, Task), f"added task must be a Task, not {type(task)}"
 
-        task.parent_task = self
+        task.parent_task.add(self)  # add myself to set of parent tasks of `task`
         self.sub_tasks.append(task)
         self.name_sub_task_map[task.name] = task
         self.responders.append(cast(Responder, task))
@@ -226,18 +226,18 @@ class Task:
             )
         else:
             self.pending_message = msg
-            if self.pending_message is not None and self.parent_task is not None:
-                # msg may have come from parent_task, so we pretend this is from
+            if self.pending_message is not None and self.caller is not None:
+                # msg may have come from `caller`, so we pretend this is from
                 # the CURRENT task's USER entity
                 self.pending_message.metadata.sender = Entity.USER
 
-        if self.parent_task is not None and self.parent_task.logger is not None:
-            self.logger = self.parent_task.logger
+        if self.caller is not None and self.caller.logger is not None:
+            self.logger = self.caller.logger
         else:
             self.logger = RichFileLogger(f"logs/{self.name}.log", color=self.color_log)
 
-        if self.parent_task is not None and self.parent_task.tsv_logger is not None:
-            self.tsv_logger = self.parent_task.tsv_logger
+        if self.caller is not None and self.caller.tsv_logger is not None:
+            self.tsv_logger = self.caller.tsv_logger
         else:
             self.tsv_logger = setup_file_logger("tsv_logger", f"logs/{self.name}.tsv")
             header = ChatDocLoggerFields().tsv_header()
@@ -250,6 +250,7 @@ class Task:
         self,
         msg: Optional[str | ChatDocument] = None,
         turns: int = -1,
+        caller: None | Task = None,
     ) -> Optional[ChatDocument]:
         """
         Loop over `step()` until task is considered done or `turns` is reached.
@@ -264,6 +265,7 @@ class Task:
                 LLM or Human (User).
             turns (int): number of turns to run the task for;
                 default is -1, which means run until task is done.
+            caller (Task|None): the calling task, if any
 
         Returns:
             Optional[ChatDocument]: valid response from the agent
@@ -281,7 +283,7 @@ class Task:
         ):
             # this task is not the intended recipient so return None
             return None
-
+        self.caller = caller
         self.init(msg)
         # sets indentation to be printed prior to any output from agent
         self.agent.indent = self._indent
@@ -447,20 +449,22 @@ class Task:
 
     def response(self, e: Responder, turns: int = -1) -> Optional[ChatDocument]:
         """
-        Get response to `self.pending_message` from an entity.
+        Get response to `self.pending_message` from a responder.
         If response is __valid__ (i.e. it ends the current turn of seeking
         responses):
             -then return the response as a ChatDocument object,
             -otherwise return None.
         Args:
-            e (Entity): entity to get response from
+            e (Responder): responder to get response from.
+            turns (int): number of turns to run the task for.
+                Default is -1, which means run until task is done.
         Returns:
             Optional[ChatDocument]: response to `self.pending_message` from entity if
             valid, None otherwise
         """
         if isinstance(e, Task):
             actual_turns = e.turns if e.turns > 0 else turns
-            return e.run(self.pending_message, turns=actual_turns)
+            return e.run(self.pending_message, turns=actual_turns, caller=self)
         else:
             return self._entity_responder_map[cast(Entity, e)](self.pending_message)
 
@@ -521,10 +525,10 @@ class Task:
             self.pending_message is None
             # LLM decided task is done
             or DONE in self.pending_message.content
-            or (  # current task is addressing message to parent task
-                self.parent_task is not None
-                and self.parent_task.name != ""
-                and self.pending_message.metadata.recipient == self.parent_task.name
+            or (  # current task is addressing message to caller task
+                self.caller is not None
+                and self.caller.name != ""
+                and self.pending_message.metadata.recipient == self.caller.name
             )
             or (
                 # Task controller is "stuck", has nothing to say

langroid/io/refs.md

@@ -0,0 +1 @@
+https://chat.openai.com/share/7c440b3f-ddbf-4ae6-a26f-ac28d947d403

langroid/language_models/azure_openai.py

@@ -0,0 +1,72 @@
+import os
+
+import openai
+from dotenv import load_dotenv
+
+from langroid.language_models.openai_gpt import OpenAIGPT, OpenAIGPTConfig
+
+
+class AzureConfig(OpenAIGPTConfig):
+    """
+    Configuration for Azure OpenAI GPT. You need to supply the env vars listed in
+    ``.azure_env_template`` after renaming the file to ``.azure_env``. Because this file
+    is used by this class to find the env vars.
+    Attributes:
+        type (str): should be ``azure``
+        api_version (str): can be set inside the ``.azure_env``
+        deployment_name (str): can be set inside the ``.azure_env`` and should be based
+        the custom name you chose for your deployment when you deployed a model
+    """
+
+    type: str = "azure"
+    api_version: str = "2023-07-01-preview"
+    deployment_name: str = ""
+
+
+class AzureGPT(OpenAIGPT):
+    """
+    Class to access OpenAI LLMs via Azure. These env variables can be obtained from the
+    file `.azure_env`. Azure OpenAI doesn't support ``completion``
+    Attributes:
+        config: AzureConfig object
+        api_key: Azure API key
+        api_base: Azure API base url
+        api_version: Azure API version
+    """
+
+    def __init__(self, config: AzureConfig):
+        super().__init__(config)
+        self.config: AzureConfig = config
+        self.api_type = config.type
+        openai.api_type = self.api_type
+        load_dotenv(dotenv_path=".azure_env")
+        self.api_key = os.getenv("AZURE_API_KEY", "")
+        if self.api_key == "":
+            raise ValueError(
+                """
+                AZURE_API_KEY not set in .env file,
+                please set it to your Azure API key."""
+            )
+
+        self.api_base = os.getenv("OPENAI_API_BASE", "")
+        if self.api_base == "":
+            raise ValueError(
+                """
+                OPENAI_API_BASE not set in .env file,
+                please set it to your Azure API key."""
+            )
+        # we don't need this for ``api_key`` because it's handled inside
+        # ``openai_gpt.py`` methods before invoking chat/completion calls
+        else:
+            openai.api_base = self.api_base
+
+        self.api_version = os.getenv("OPENAI_API_VERSION", "") or config.api_version
+        openai.api_version = self.api_version
+
+        self.deployment_name = os.getenv("OPENAI_DEPLOYMENT_NAME", "")
+        if self.deployment_name == "":
+            raise ValueError(
+                """
+                OPENAI_DEPLOYMENT_NAME not set in .env file,
+                please set it to your Azure API key."""
+            )

langroid/language_models/base.py

@@ -36,6 +36,9 @@ class LLMConfig(BaseSettings):
     stream: bool = False  # stream output from API?
     cache_config: None | RedisCacheConfig | MomentoCacheConfig = None
 
+    # Dict of model -> (input/prompt cost, output/completion cost)
+    cost_per_1k_tokens: Optional[Dict[str, Tuple[float, float]]] = None
+
 
 class LLMFunctionCall(BaseModel):
     """
@@ -63,6 +66,16 @@ class LLMFunctionSpec(BaseModel):
     parameters: Dict[str, Any]
 
 
+class LLMTokenUsage(BaseModel):
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    cost: float = 0.0
+
+    @property
+    def total_tokens(self) -> int:
+        return self.prompt_tokens + self.completion_tokens
+
+
 class Role(str, Enum):
     USER = "user"
     SYSTEM = "system"
@@ -116,7 +129,7 @@ class LLMResponse(BaseModel):
 
     message: str
     function_call: Optional[LLMFunctionCall] = None
-    usage: int
+    usage: Optional[LLMTokenUsage]
     cached: bool = False
 
     def to_LLMMessage(self) -> LLMMessage:
@@ -193,13 +206,21 @@ class LanguageModel(ABC):
             config: configuration for language model
         Returns: instance of language model
         """
+        from langroid.language_models.azure_openai import AzureGPT
         from langroid.language_models.openai_gpt import OpenAIGPT
 
         if config is None or config.type is None:
             return None
+
+        openai: Union[Type[AzureGPT], Type[OpenAIGPT]]
+
+        if config.type == "azure":
+            openai = AzureGPT
+        else:
+            openai = OpenAIGPT
         cls = dict(
-            openai=OpenAIGPT,
-        ).get(config.type, OpenAIGPT)
+            openai=openai,
+        ).get(config.type, openai)
         return cls(config)  # type: ignore
 
     @abstractmethod
@@ -248,6 +269,13 @@ class LanguageModel(ABC):
             raise ValueError("No context length  specified")
         return self.config.context_length[self.config.completion_model]
 
+    def chat_cost(self) -> Tuple[float, float]:
+        if self.config.chat_model is None:
+            raise ValueError("No chat model specified")
+        if self.config.cost_per_1k_tokens is None:
+            raise ValueError("No cost per 1k tokens  specified")
+        return self.config.cost_per_1k_tokens[self.config.chat_model]
+
     def followup_to_standalone(
         self, chat_history: List[Tuple[str, str]], question: str
     ) -> str:
@@ -368,7 +396,10 @@ class LanguageModel(ABC):
             sources = ""
         return Document(
             content=content,
-            metadata={"source": "SOURCE: " + sources, "cached": llm_response.cached},
+            metadata={
+                "source": "SOURCE: " + sources,
+                "cached": llm_response.cached,
+            },
         )
 
 

langroid/language_models/openai_gpt.py

@@ -20,6 +20,7 @@ from langroid.language_models.base import (
     LLMFunctionSpec,
     LLMMessage,
     LLMResponse,
+    LLMTokenUsage,
     Role,
 )
 from langroid.language_models.utils import (
@@ -62,6 +63,12 @@ class OpenAIGPTConfig(LLMConfig):
         OpenAIChatModel.GPT4_NOFUNC: 8192,
         OpenAICompletionModel.TEXT_DA_VINCI_003: 4096,
     }
+    cost_per_1k_tokens: Dict[str, Tuple[float, float]] = {
+        # (input/prompt cost, output/completion cost)
+        OpenAIChatModel.GPT3_5_TURBO: (0.0015, 0.002),
+        OpenAIChatModel.GPT4: (0.03, 0.06),  # 8K context
+        OpenAIChatModel.GPT4_NOFUNC: (0.03, 0.06),
+    }
 
 
 class OpenAIResponse(BaseModel):
@@ -208,7 +215,6 @@ class OpenAIGPT(LanguageModel):
         return (  # type: ignore
             LLMResponse(
                 message=completion,
-                usage=0,
                 cached=False,
                 function_call=function_call if has_function else None,
             ),
@@ -229,6 +235,31 @@ class OpenAIGPT(LanguageModel):
         # Try to get the result from the cache
         return hashed_key, self.cache.retrieve(hashed_key)
 
+    def _cost_chat_model(self, prompt: int, completion: int) -> float:
+        price = self.chat_cost()
+        return (price[0] * prompt + price[1] * completion) / 1000
+
+    def _handle_token_usage(
+        self, cached: bool, response: Dict[str, Any]
+    ) -> LLMTokenUsage:
+        cost = 0.0
+        prompt_tokens = 0
+        completion_tokens = 0
+        if not cached and not self.config.stream:
+            prompt_tokens = response["usage"]["prompt_tokens"]
+            completion_tokens = response["usage"]["completion_tokens"]
+            cost = self._cost_chat_model(
+                response["usage"]["prompt_tokens"],
+                response["usage"]["completion_tokens"],
+            )
+        # if not self.config.stream:
+        #     prompt_tokens = response["usage"]["prompt_tokens"]
+        #     completion_tokens = response["usage"]["completion_tokens"]
+
+        return LLMTokenUsage(
+            prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, cost=cost
+        )
+
     def generate(self, prompt: str, max_tokens: int) -> LLMResponse:
         try:
             return self._generate(prompt, max_tokens)
@@ -236,7 +267,7 @@ class OpenAIGPT(LanguageModel):
             # capture exceptions not handled by retry, so we don't crash
             err_msg = str(e)[:500]
             logging.error(f"OpenAI API error: {err_msg}")
-            return LLMResponse(message=NO_ANSWER, usage=0, cached=False)
+            return LLMResponse(message=NO_ANSWER, cached=False)
 
     def _generate(self, prompt: str, max_tokens: int) -> LLMResponse:
         if self.config.use_chat_for_completion:
@@ -265,8 +296,9 @@ class OpenAIGPT(LanguageModel):
                     self.cache.store(hashed_key, result)
             return cached, hashed_key, result
 
+        key_name = "engine" if self.config.type == "azure" else "model"
         cached, hashed_key, response = completions_with_backoff(
-            model=self.config.completion_model,
+            **{key_name: self.config.completion_model},
             prompt=prompt,
             max_tokens=max_tokens,  # for output/completion
             request_timeout=self.config.timeout,
@@ -275,9 +307,8 @@ class OpenAIGPT(LanguageModel):
             stream=self.config.stream,
         )
 
-        usage = response["usage"]["total_tokens"]
         msg = response["choices"][0]["text"].strip()
-        return LLMResponse(message=msg, usage=usage, cached=cached)
+        return LLMResponse(message=msg, cached=cached)
 
     async def agenerate(self, prompt: str, max_tokens: int) -> LLMResponse:
         try:
@@ -286,7 +317,7 @@ class OpenAIGPT(LanguageModel):
             # capture exceptions not handled by retry, so we don't crash
             err_msg = str(e)[:500]
             logging.error(f"OpenAI API error: {err_msg}")
-            return LLMResponse(message=NO_ANSWER, usage=0, cached=False)
+            return LLMResponse(message=NO_ANSWER, cached=False)
 
     async def _agenerate(self, prompt: str, max_tokens: int) -> LLMResponse:
         openai.api_key = self.api_key
@@ -324,7 +355,6 @@ class OpenAIGPT(LanguageModel):
                 temperature=self.config.temperature,
                 stream=self.config.stream,
             )
-            usage = response["usage"]["total_tokens"]
             msg = response["choices"][0]["message"]["content"].strip()
         else:
 
@@ -349,9 +379,8 @@ class OpenAIGPT(LanguageModel):
                 echo=False,
                 stream=self.config.stream,
             )
-            usage = response["usage"]["total_tokens"]
             msg = response["choices"][0]["text"].strip()
-        return LLMResponse(message=msg, usage=usage, cached=cached)
+        return LLMResponse(message=msg, cached=cached)
 
     def chat(
         self,
@@ -366,7 +395,7 @@ class OpenAIGPT(LanguageModel):
             # capture exceptions not handled by retry, so we don't crash
             err_msg = str(e)[:500]
             logging.error(f"OpenAI API error: {err_msg}")
-            return LLMResponse(message=NO_ANSWER, usage=0, cached=False)
+            return LLMResponse(message=NO_ANSWER, cached=False)
 
     def _chat(
         self,
@@ -421,8 +450,14 @@ class OpenAIGPT(LanguageModel):
                     self.cache.store(hashed_key, result)
             return cached, hashed_key, result
 
+        if self.config.type == "azure":
+            key_name = "engine"
+            if hasattr(self, "deployment_name"):
+                self.config.chat_model = self.deployment_name
+        else:
+            key_name = "model"
         args: Dict[str, Any] = dict(
-            model=self.config.chat_model,
+            **{key_name: self.config.chat_model},
             messages=[m.api_dict() for m in llm_messages],
             max_tokens=max_tokens,
             n=1,
@@ -447,7 +482,6 @@ class OpenAIGPT(LanguageModel):
             self.cache.store(hashed_key, openai_response)
             return llm_response
 
-        usage = response["usage"]["total_tokens"]
         # openAI response will look like this:
         """
         {
@@ -477,7 +511,6 @@ class OpenAIGPT(LanguageModel):
             }
         }
         """
-
         message = response["choices"][0]["message"]
         msg = message["content"] or ""
         if message.get("function_call") is None:
@@ -489,10 +522,10 @@ class OpenAIGPT(LanguageModel):
                 fun_call.arguments = fun_args
             except (ValueError, SyntaxError):
                 logging.warning(
-                    f"Could not parse function arguments: "
+                    "Could not parse function arguments: "
                     f"{message['function_call']['arguments']} "
                     f"for function {message['function_call']['name']} "
-                    f"treating as normal non-function message"
+                    "treating as normal non-function message"
                 )
                 fun_call = None
                 msg = message["content"] + message["function_call"]["arguments"]
@@ -500,6 +533,6 @@ class OpenAIGPT(LanguageModel):
         return LLMResponse(
             message=msg.strip() if msg is not None else "",
             function_call=fun_call,
-            usage=usage,
             cached=cached,
+            usage=self._handle_token_usage(cached, response),
         )

langroid/prompts/templates.py

@@ -8,15 +8,16 @@ EXTRACT_RELEVANT = """
     Question:{question}
     Relevant text, if any: """.strip()
 
-EXTRACTION_PROMPT_GPT4 = """
+EXTRACTION_PROMPT_GPT4 = f"""
 Given the content and question below, extract COMPLETE SENTENCES OR PHRASES 
 VERBATIM from the content, that are relevant to answering the question (if such text 
 exists), even if it contradicts your knowledge, and even if it is factually incorrect.
 Do not  make up an answer that is not supported by the content. 
-When you answer, be concise, no need to explain anything. 
+When you answer, be concise, no need to explain anything. If there is no relevant text,
+simply say {NO_ANSWER}.
 
-Content: {content}
-Question: {question}
+Content: {{content}}
+Question: {{question}}
 Relevant text, if any:
 """
 

{langroid-0.1.52.dist-info → langroid-0.1.54.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langroid
-Version: 0.1.52
+Version: 0.1.54
 Summary: Harness LLMs with Multi-Agent Programming
 License: MIT
 Author: Prasad Chalasani
@@ -75,7 +75,7 @@ Description-Content-Type: text/markdown
 
 <div align="center">
 
-[![PyPI version](https://badge.fury.io/py/langroid.svg)](https://badge.fury.io/py/langroid)
+[![PyPI - Version](https://img.shields.io/pypi/v/langroid)](https://pypi.org/project/langroid/)
 [![Pytest](https://github.com/langroid/langroid/actions/workflows/pytest.yml/badge.svg)](https://github.com/langroid/langroid/actions/workflows/pytest.yml)
 [![codecov](https://codecov.io/gh/langroid/langroid/branch/main/graph/badge.svg?token=H94BX5F0TE)](https://codecov.io/gh/langroid/langroid)
 [![Lint](https://github.com/langroid/langroid/actions/workflows/validate.yml/badge.svg)](https://github.com/langroid/langroid/actions/workflows/validate.yml)
@@ -135,6 +135,7 @@ for ideas on what to contribute.
 <summary> <b>:fire: Updates/Releases</b></summary>
 
 - **Aug 2023:**
+  - **[Hierarchical computation](https://langroid.github.io/langroid/examples/agent-tree/)** example using Langroid agents and task orchestration.
   - **0.1.51:** Support for global state, see [test_global_state.py](tests/main/test_global_state.py).
   - **:whale: Langroid Docker image**, available, see instructions below.
   - [**RecipientTool**](langroid/agent/tools/recipient_tool.py) enables (+ enforces) LLM to 
@@ -328,6 +329,26 @@ GOOGLE_CSE_ID=your-cse-id
 ```
 </details>
 
+<details>
+<summary><b>Setup instructions for Microsoft Azure OpenAI(click to expand)</b></summary> 
+In the root of the repo, copy the `.azure_env_template` file to a new file `.azure_env`: 
+
+```bash
+cp .azure_env_template .azure_env
+```
+
+The file `.azure_env` contains four environment variables that are required to use Azure OpenAI: `AZURE_API_KEY`, `OPENAI_API_BASE`, `OPENAI_API_VERSION`, and `OPENAI_DEPLOYMENT_NAME`
+
+This page [Microsoft Azure OpenAI](https://learn.microsoft.com/en-us/azure/ai-services/openai/chatgpt-quickstart?tabs=command-line&pivots=programming-language-python#environment-variables)
+provides more information, and you can set each environment variable as follows:
+
+- `AZURE_API_KEY`, from the value of `API_KEY`
+- `OPENAI_API_BASE` from the value of `ENDPOINT`, typically looks like `https://your.domain.azure.com`.
+- For `OPENAI_API_VERSION`, you can use the default value in `.azure_env_template`, and latest version can be found [here](https://learn.microsoft.com/en-us/azure/ai-services/openai/whats-new#azure-openai-chat-completion-general-availability-ga)
+- `OPENAI_DEPLOYMENT_NAME` is the deployment name you chose when you deployed the GPT-35-Turbo or GPT-4 models.
+
+</details>
+
 ---
 
 # :whale: Docker Instructions

{langroid-0.1.52.dist-info → langroid-0.1.54.dist-info}/RECORD

@@ -1,12 +1,12 @@
 langroid/__init__.py,sha256=sEKJ_5WJBAMZApevfeE3gxLK-eotVzJMJlT83G0rAko,30
 langroid/agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langroid/agent/base.py,sha256=_qlyXDf4MOZ7ZyPlaVyDIMPyE4gO2huf-rUJI6FgodE,23717
-langroid/agent/chat_agent.py,sha256=SAU06R5kv6bisQ8B5UdkMdHa89Xji6GvMrkjx9Ly0Uo,22529
-langroid/agent/chat_document.py,sha256=bh1xKUcTkikSrNbJ1mm0NP1aGVvJZr94LMZvEPM2nY4,6072
+langroid/agent/base.py,sha256=bnqa_PZsw1_RWDv1w67g1rMrhbGTdt_mTPWcZ_uAZIk,26530
+langroid/agent/chat_agent.py,sha256=Sma0-5XPHDzBOcduthwwlWBmkBgqpk8gGzStF8rcrps,22643
+langroid/agent/chat_document.py,sha256=apaYj38sDu7ALCnsA8tJwoj3Z8zLNmIsNPd4-IujnGk,6153
 langroid/agent/helpers.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/agent/junk,sha256=LxfuuW7Cijsg0szAzT81OjWWv1PMNI-6w_-DspVIO2s,339
 langroid/agent/special/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langroid/agent/special/doc_chat_agent.py,sha256=KveFWGpiz3Md0_Zy_Q_AdLG-5rHl5wSparwhKOD2tOE,16486
+langroid/agent/special/doc_chat_agent.py,sha256=yyf_kXj0rFmNWY_TjllGstZBg_hfFNDLiZ4p5P_P5Rg,16528
 langroid/agent/special/recipient_validator_agent.py,sha256=R3Rit93BNWQar_9stuDBGzmLr2W-IYOQ7oq-tlNNlps,6035
 langroid/agent/special/retriever_agent.py,sha256=DeOB5crFjXBvDEZT9k9ZVinOfFM2VgS6tQWWFyXSk9o,7204
 langroid/agent/special/sql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -14,7 +14,7 @@ langroid/agent/special/sql/sql_chat_agent.py,sha256=DVvBI_q1HqTAk_r_4-B2vGOWjneb
 langroid/agent/special/sql/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/agent/special/sql/utils/description_extractors.py,sha256=mxgbIZMqGnMj3vtU614bJJSpXY-JLjhX4iEB2R1-qlI,3921
 langroid/agent/special/table_chat_agent.py,sha256=PTCE7MmunQj7tFiKAMIh7kvdIeQYU5ceXgBabwsxdg8,7244
-langroid/agent/task.py,sha256=sRm_2w3ETPlXvu_Zo6_R6AU6oYvn_nhdm0aB3Tea3T0,27783
+langroid/agent/task.py,sha256=de8JY1Nq3YZ13dThPSYoVgj7BHxbXyifHiQtKQi5vQc,28111
 langroid/agent/tool_message.py,sha256=08vKug6t6m8aSNuO268bc3xfz-rivDU5brxBOSGdqNM,5935
 langroid/agent/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/agent/tools/google_search_tool.py,sha256=64F9oMNdS237BBOitrvYXN4Il_ES_fNrHkh35tBEDfA,1160
@@ -28,9 +28,14 @@ langroid/embedding_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJW
 langroid/embedding_models/base.py,sha256=176jDrjEAAhNzdFCG8pfossd8SAhvHR8Q5Y8pOOm0LI,983
 langroid/embedding_models/clustering.py,sha256=tZWElUqXl9Etqla0FAa7og96iDKgjqWjucZR_Egtp-A,6684
 langroid/embedding_models/models.py,sha256=1xcv9hqmCTsbUbS8v7XeZRsf25Tu79JUoSipIYpvNoo,2765
+langroid/io/base.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langroid/io/cmd_line.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+langroid/io/refs.md,sha256=B1lrcHhg0IOWm74OHHZrv8o41pVvX6hSiG5pJjpDF90,67
+langroid/io/websocket.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/language_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-langroid/language_models/base.py,sha256=B11jZ6fLSbg1rS1AXM0i_McLeeKpmXmR9vwjY2oOiKc,13751
-langroid/language_models/openai_gpt.py,sha256=eeX3gUFzQxUS2v-O5vMrLqiMO1-_T8mqM4mtVerL_UM,19377
+langroid/language_models/azure_openai.py,sha256=Axzg-oysiMwKgJR6uGLBcnqVFf3NgykVN0b0Mex2rGY,2543
+langroid/language_models/base.py,sha256=CHSMWJd9kFwMsI38pLmFcPtgkBUUQ3a47sj77kD8-bw,14743
+langroid/language_models/openai_gpt.py,sha256=-qiQV2OldX9PShHX7UqknGTnudPpCn2C7n1NL2S-Be4,20748
 langroid/language_models/utils.py,sha256=rmnSn-sJ3aKl_wBdeLPkck0Li4Ed6zkCxZYYl7n1V34,4668
 langroid/mytypes.py,sha256=YA42IJcooJnTxAwk-B4FmZ1hqzIIF1ZZKcpUKzBTGGo,1537
 langroid/parsing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -51,7 +56,7 @@ langroid/parsing/web_search.py,sha256=hGUVoSJNdpoT5rsm-ikAteMiUropHrzKaxN8EVVqO2
 langroid/prompts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/prompts/dialog.py,sha256=SpfiSyofSgy2pwD1YboHR_yHO3LEEMbv6j2sm874jKo,331
 langroid/prompts/prompts_config.py,sha256=EMK1Fm7EmS8y3CV4AkrVgn5K4NipiM4m7J8819W1KeM,98
-langroid/prompts/templates.py,sha256=4OAujKJzKIhqt4SQL2GE1aPahlZuDotlT_dZAKx0bqE,6311
+langroid/prompts/templates.py,sha256=4X-07tnmUQ8Z_zaWRQAUUyKiErGztp3tERujqnG8sGA,6369
 langroid/prompts/transforms.py,sha256=GsQo1klGxUy0fACh6j0lTblk6XEl2erRnhRWlN2M4-c,2706
 langroid/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -74,7 +79,7 @@ langroid/vector_store/base.py,sha256=QZx3NUNwf2I0r3A7iuoUHIRGbqt_pFGD0hq1R-Yg8iM
 langroid/vector_store/chromadb.py,sha256=s5pQkKjaMP-Tt5A8M10EInFzttaALPbJAq7q4gf0TKg,5235
 langroid/vector_store/qdrant_cloud.py,sha256=3im4Mip0QXLkR6wiqVsjV1QvhSElfxdFSuDKddBDQ-4,188
 langroid/vector_store/qdrantdb.py,sha256=KRvIIj1IZG2zFqejofMnRs2hT86B-27LgBEnuczdqOU,9072
-langroid-0.1.52.dist-info/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
-langroid-0.1.52.dist-info/WHEEL,sha256=vVCvjcmxuUltf8cYhJ0sJMRDLr1XsPuxEId8YDzbyCY,88
-langroid-0.1.52.dist-info/METADATA,sha256=dEWMjtwq9vm5rq8-WyzkTJZGeNT6k-2WORNeopAoIi4,33607
-langroid-0.1.52.dist-info/RECORD,,
+langroid-0.1.54.dist-info/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
+langroid-0.1.54.dist-info/WHEEL,sha256=vVCvjcmxuUltf8cYhJ0sJMRDLr1XsPuxEId8YDzbyCY,88
+langroid-0.1.54.dist-info/METADATA,sha256=ZNXtQxEqxLa1l-4Bd2WHksnrCwAuv2Xf3zpapwZUahM,34974
+langroid-0.1.54.dist-info/RECORD,,