add langchain llm math chain

evals/completion_fns/langchain_math.py

@@ -0,0 +1,30 @@
+import importlib
+from typing import Optional
+
+from langchain import OpenAI, LLMMathChain
+
+from evals.prompt.base import CompletionPrompt
+from evals.record import record_sampling
+
+
+class LangChainCompletionResult:
+    def __init__(self, response) -> None:
+        self.response = response
+
+    def get_completions(self) -> list[str]:
+        return [self.response.strip()]
+
+
+class LangChainMathChainCompletionFn:
+    def __init__(self, **kwargs) -> None:
+        llm = OpenAI(temperature=0)
+        self.llm_math = LLMMathChain(llm=llm)
+
+    def __call__(self, prompt, **kwargs) -> LangChainCompletionResult:
+
+        prompt = CompletionPrompt(prompt).to_formatted_prompt()
+        response = self.llm_math.run(prompt)
+        # The LangChain response comes with `Answer: ` ahead of this, let's strip it out
+        response = response.strip("Answer:").strip()
+        record_sampling(prompt=prompt, sampled=response)
+        return LangChainCompletionResult(response)

evals/registry.py

@@ -104,7 +104,9 @@ def make_completion_fn(self, name: str) -> CompletionFn:
         spec = self.get_completion_fn(name)
         if spec is None:
             raise ValueError(f"Could not find CompletionFn in the registry with ID {name}")
-
+        if spec.args is None:
+            spec.args = {}
+
         spec.args["registry"] = self
         instance = make_object(spec.cls)(**spec.args or {})
         assert isinstance(instance, CompletionFn), f"{name} must be a CompletionFn"

evals/registry/completion_fns/langchain_chains.yaml

@@ -0,0 +1,2 @@
+langchain/chains/llm_math:
+  class: evals.completion_fns.langchain_math:LangChainMathChainCompletionFn