langfun 0.1.2.dev202412090805__py3-none-any.whl → 0.1.2.dev202412130804__py3-none-any.whl

langfun/core/agentic/action_test.py

@@ -18,6 +18,8 @@ import unittest
 import langfun.core as lf
 from langfun.core.agentic import action as action_lib
 from langfun.core.llms import fake
+import langfun.core.structured as lf_structured
+import pyglove as pg
 
 
 class SessionTest(unittest.TestCase):
@@ -28,60 +30,94 @@ class SessionTest(unittest.TestCase):
     class Bar(action_lib.Action):
 
       def call(self, session, *, lm, **kwargs):
-        test.assertIs(session.current_invocation.action, self)
+        test.assertIs(session.current_action.action, self)
         session.info('Begin Bar')
         session.query('bar', lm=lm)
+        session.add_metadata(note='bar')
         return 2
 
     class Foo(action_lib.Action):
       x: int
 
       def call(self, session, *, lm, **kwargs):
-        test.assertIs(session.current_invocation.action, self)
-        session.info('Begin Foo', x=1)
-        session.query('foo', lm=lm)
+        test.assertIs(session.current_action.action, self)
+        with session.phase('prepare'):
+          session.info('Begin Foo', x=1)
+          session.query('foo', lm=lm)
+        with session.track_queries():
+          self.make_additional_query(lm)
+        session.add_metadata(note='foo')
         return self.x + Bar()(session, lm=lm)
 
+      def make_additional_query(self, lm):
+        lf_structured.query('additional query', lm=lm)
+
     lm = fake.StaticResponse('lm response')
-    session = action_lib.Session()
-    root = session.root_invocation
-    self.assertIsInstance(root.action, action_lib.RootAction)
-    self.assertIs(session.current_invocation, session.root_invocation)
-    self.assertEqual(Foo(1)(session, lm=lm), 3)
-    self.assertEqual(len(session.root_invocation.child_invocations), 1)
-    self.assertEqual(len(list(session.root_invocation.queries())), 0)
-    self.assertEqual(
-        len(list(session.root_invocation.queries(include_children=True))), 2
-    )
-    self.assertEqual(
-        len(list(session.root_invocation.child_invocations[0].queries())), 1
-    )
-    self.assertEqual(len(session.root_invocation.child_invocations[0].logs), 1)
-    self.assertEqual(
-        len(session.root_invocation.child_invocations[0].child_invocations),
-        1
-    )
-    self.assertEqual(
-        len(session.root_invocation
-            .child_invocations[0].child_invocations[0].logs),
-        1
-    )
-    self.assertEqual(
-        len(list(session.root_invocation
-                 .child_invocations[0].child_invocations[0].queries())),
-        1
-    )
-    self.assertEqual(
-        len(session.root_invocation
-            .child_invocations[0].child_invocations[0].child_invocations),
-        0
-    )
-    self.assertIs(session.current_invocation, session.root_invocation)
-    self.assertIs(session.final_result, 3)
-    self.assertIn(
-        'invocation-final-result',
-        session.to_html().content,
+    foo = Foo(1)
+    self.assertEqual(foo(lm=lm), 3)
+
+    session = foo.session
+    self.assertIsNotNone(session)
+    self.assertIsInstance(session.root.action, action_lib.RootAction)
+    self.assertIs(session.current_action, session.root)
+
+    #
+    # Inspecting the root invocation.
+    #
+
+    root = session.root
+    self.assertEqual(len(root.execution.items), 1)
+    self.assertIs(root.execution.items[0].action, foo)
+
+    self.assertTrue(root.execution.has_started)
+    self.assertTrue(root.execution.has_stopped)
+    self.assertGreater(root.execution.elapse, 0)
+    self.assertEqual(root.result, 3)
+    self.assertEqual(root.metadata, dict(note='foo'))
+
+    # The root space should have one action (foo), no queries, and no logs.
+    self.assertEqual(len(list(root.actions)), 1)
+    self.assertEqual(len(list(root.queries)), 0)
+    self.assertEqual(len(list(root.logs)), 0)
+    # 1 query from Bar and 2 from Foo.
+    self.assertEqual(len(list(root.all_queries)), 3)
+    # 1 log from Bar and 1 from Foo.
+    self.assertEqual(len(list(root.all_logs)), 2)
+    self.assertEqual(root.usage_summary.total.num_requests, 3)
+
+    # Inspecting the top-level action (Foo)
+    foo_invocation = root.execution.items[0]
+    self.assertEqual(len(foo_invocation.execution.items), 3)
+
+    # Prepare phase.
+    prepare_phase = foo_invocation.execution.items[0]
+    self.assertIsInstance(
+        prepare_phase, action_lib.ExecutionTrace
     )
+    self.assertEqual(len(prepare_phase.items), 2)
+    self.assertTrue(prepare_phase.has_started)
+    self.assertTrue(prepare_phase.has_stopped)
+    self.assertEqual(prepare_phase.usage_summary.total.num_requests, 1)
+
+    # Tracked queries.
+    query_invocation = foo_invocation.execution.items[1]
+    self.assertIsInstance(query_invocation, lf_structured.QueryInvocation)
+    self.assertIs(query_invocation.lm, lm)
+
+    # Invocation to Bar.
+    bar_invocation = foo_invocation.execution.items[2]
+    self.assertIsInstance(bar_invocation, action_lib.ActionInvocation)
+    self.assertIsInstance(bar_invocation.action, Bar)
+    self.assertEqual(bar_invocation.result, 2)
+    self.assertEqual(bar_invocation.metadata, dict(note='bar'))
+    self.assertEqual(len(bar_invocation.execution.items), 2)
+
+    # Save to HTML
+    self.assertIn('result', session.to_html().content)
+
+    # Save session to JSON
+    json_str = session.to_json_str(save_ref_value=True)
+    self.assertIsInstance(pg.from_json_str(json_str), action_lib.Session)
 
   def test_log(self):
     session = action_lib.Session()

langfun/core/eval/v2/evaluation.py

@@ -556,6 +556,9 @@ class Evaluation(experiment_lib.Experiment):
           border: 0px;
           margin: 0px;
         }
+        .eval-details .tab-control {
+          width: 100%;
+        }
         .eval-details .tab-button {
           font-size: large;
         }

langfun/core/eval/v2/experiment.py

@@ -465,6 +465,33 @@ class Experiment(lf.Component, pg.views.HtmlTreeView.Extension):
     runner.run()
     return runner.current_run
 
+  def run_preconfigured(
+      self,
+      root_dir: str | None = None,
+      id: str | None = None,   # pylint: disable=redefined-builtin
+      **kwargs
+  ) -> 'Run':
+    """Runs the experiment with pre-configured kwargs from `cls.RUN_ARGS`.
+
+    This helper method allows users to config running arguments as a part of
+    the class.
+
+    Args:
+      root_dir: root directory of the experiment.
+      id: ID of the current run.
+      **kwargs: Keyword arguments to override the RUN_CONFIG.
+
+    Returns:
+      The current run.
+    """
+    run_config = getattr(self, 'RUN_ARGS', {})
+    run_config.update(kwargs)
+    if root_dir is not None:
+      run_config['root_dir'] = root_dir
+    if id is not None:
+      run_config['id'] = id
+    return self.run(**run_config)
+
   #
   # HTML views.
   #

langfun/core/eval/v2/experiment_test.py

@@ -39,6 +39,10 @@ def sample_inputs():
 
 class MyEvaluation(Evaluation):
   NAME = 'my_eval'
+  RUN_ARGS = dict(
+      runner='test'
+  )
+
   replica_id: int = 0
   inputs = sample_inputs()
   metrics = [metrics_lib.Match()]
@@ -288,10 +292,17 @@ class RunnerTest(unittest.TestCase):
         TestRunner
     )
     root_dir = os.path.join(tempfile.gettempdir(), 'my_eval')
+
+    # Test standard run.
     MyEvaluation(replica_id=0).run(
         root_dir, id='20241101_0', runner='test'
     )
 
+    # Test run preconfigured.
+    MyEvaluation(replica_id=0).run_preconfigured(
+        root_dir=root_dir, id='20241101_1'
+    )
+
     with self.assertRaisesRegex(
         ValueError, 'Runner class must define a NAME constant'
     ):

langfun/core/llms/__init__.py

@@ -32,10 +32,12 @@ from langfun.core.llms.rest import REST
 
 # Gemini models.
 from langfun.core.llms.google_genai import GenAI
+from langfun.core.llms.google_genai import GeminiFlash2_0Exp
 from langfun.core.llms.google_genai import GeminiExp_20241114
+from langfun.core.llms.google_genai import GeminiExp_20241206
 from langfun.core.llms.google_genai import GeminiFlash1_5
-from langfun.core.llms.google_genai import GeminiPro
 from langfun.core.llms.google_genai import GeminiPro1_5
+from langfun.core.llms.google_genai import GeminiPro
 from langfun.core.llms.google_genai import GeminiProVision
 from langfun.core.llms.google_genai import Palm2
 from langfun.core.llms.google_genai import Palm2_IT
@@ -120,6 +122,8 @@ from langfun.core.llms.groq import GroqWhisper_Large_v3
 from langfun.core.llms.groq import GroqWhisper_Large_v3Turbo
 
 from langfun.core.llms.vertexai import VertexAI
+from langfun.core.llms.vertexai import VertexAIGemini2_0
+from langfun.core.llms.vertexai import VertexAIGeminiFlash2_0Exp
 from langfun.core.llms.vertexai import VertexAIGemini1_5
 from langfun.core.llms.vertexai import VertexAIGeminiPro1_5
 from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_001

langfun/core/llms/anthropic.py

@@ -35,7 +35,7 @@ except ImportError:
   google_auth = None
   auth_requests = None
   credentials_lib = None
-  Credentials = None  # pylint: disable=invalid-name
+  Credentials = Any  # pylint: disable=invalid-name
 
 
 SUPPORTED_MODELS_AND_SETTINGS = {

langfun/core/llms/google_genai.py

@@ -20,6 +20,7 @@ from typing import Annotated, Any, Literal
 
 import langfun.core as lf
 from langfun.core import modalities as lf_modalities
+from langfun.core.llms import vertexai
 import pyglove as pg
 
 
@@ -47,13 +48,15 @@ class GenAI(lf.LanguageModel):
 
   model: Annotated[
       Literal[
+          'gemini-2.0-flash-exp',
+          'gemini-exp-1206',
+          'gemini-exp-1114',
+          'gemini-1.5-pro-latest',
+          'gemini-1.5-flash-latest',
           'gemini-pro',
           'gemini-pro-vision',
           'text-bison-001',
           'chat-bison-001',
-          'gemini-1.5-pro-latest',
-          'gemini-1.5-flash-latest',
-          'gemini-exp-1114',
       ],
       'Model name.',
   ]
@@ -306,64 +309,64 @@ _GOOGLE_GENAI_MODEL_HUB = _ModelHub()
 #
 
 
-_IMAGE_TYPES = [
-    'image/png',
-    'image/jpeg',
-    'image/webp',
-    'image/heic',
-    'image/heif',
-]
-
-_AUDIO_TYPES = [
-    'audio/aac',
-    'audio/flac',
-    'audio/mp3',
-    'audio/m4a',
-    'audio/mpeg',
-    'audio/mpga',
-    'audio/mp4',
-    'audio/opus',
-    'audio/pcm',
-    'audio/wav',
-    'audio/webm'
-]
-
-_VIDEO_TYPES = [
-    'video/mov',
-    'video/mpeg',
-    'video/mpegps',
-    'video/mpg',
-    'video/mp4',
-    'video/webm',
-    'video/wmv',
-    'video/x-flv',
-    'video/3gpp',
-]
-
-_PDF = [
-    'application/pdf',
-]
+class GeminiFlash2_0Exp(GenAI):  # pylint: disable=invalid-name
+  """Gemini Experimental model launched on 12/06/2024."""
+
+  model = 'gemini-2.0-flash-exp'
+  supported_modalities = (
+      vertexai.DOCUMENT_TYPES
+      + vertexai.IMAGE_TYPES
+      + vertexai.AUDIO_TYPES
+      + vertexai.VIDEO_TYPES
+  )
+
+
+class GeminiExp_20241206(GenAI):  # pylint: disable=invalid-name
+  """Gemini Experimental model launched on 12/06/2024."""
+
+  model = 'gemini-exp-1206'
+  supported_modalities = (
+      vertexai.DOCUMENT_TYPES
+      + vertexai.IMAGE_TYPES
+      + vertexai.AUDIO_TYPES
+      + vertexai.VIDEO_TYPES
+  )
 
 
 class GeminiExp_20241114(GenAI):  # pylint: disable=invalid-name
   """Gemini Experimental model launched on 11/14/2024."""
 
   model = 'gemini-exp-1114'
-  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
+  supported_modalities = (
+      vertexai.DOCUMENT_TYPES
+      + vertexai.IMAGE_TYPES
+      + vertexai.AUDIO_TYPES
+      + vertexai.VIDEO_TYPES
+  )
 
 
 class GeminiPro1_5(GenAI):  # pylint: disable=invalid-name
   """Gemini Pro latest model."""
 
   model = 'gemini-1.5-pro-latest'
-  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
+  supported_modalities = (
+      vertexai.DOCUMENT_TYPES
+      + vertexai.IMAGE_TYPES
+      + vertexai.AUDIO_TYPES
+      + vertexai.VIDEO_TYPES
+  )
 
 
 class GeminiFlash1_5(GenAI):  # pylint: disable=invalid-name
   """Gemini Flash latest model."""
 
   model = 'gemini-1.5-flash-latest'
-  supported_modalities = _PDF + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
+  supported_modalities = (
+      vertexai.DOCUMENT_TYPES
+      + vertexai.IMAGE_TYPES
+      + vertexai.AUDIO_TYPES
+      + vertexai.VIDEO_TYPES
+  )
 
 
 class GeminiPro(GenAI):
@@ -376,7 +379,7 @@ class GeminiProVision(GenAI):
   """Gemini Pro vision model."""
 
   model = 'gemini-pro-vision'
-  supported_modalities = _IMAGE_TYPES + _VIDEO_TYPES
+  supported_modalities = vertexai.IMAGE_TYPES + vertexai.VIDEO_TYPES
 
 
 class Palm2(GenAI):

langfun/core/llms/vertexai.py

@@ -40,7 +40,7 @@ except ImportError:
 
 # https://cloud.google.com/vertex-ai/generative-ai/pricing
 # describes that the average number of characters per token is about 4.
-AVGERAGE_CHARS_PER_TOEKN = 4
+AVGERAGE_CHARS_PER_TOKEN = 4
 
 
 # Price in US dollars,
@@ -102,6 +102,18 @@ SUPPORTED_MODELS_AND_SETTINGS = {
         cost_per_1k_input_chars=0.000125,
         cost_per_1k_output_chars=0.000375,
     ),
+    # TODO(sharatsharat): Update costs when published
+    'gemini-exp-1206': pg.Dict(
+        rpm=20,
+        cost_per_1k_input_chars=0.000,
+        cost_per_1k_output_chars=0.000,
+    ),
+    # TODO(sharatsharat): Update costs when published
+    'gemini-2.0-flash-exp': pg.Dict(
+        rpm=20,
+        cost_per_1k_input_chars=0.000,
+        cost_per_1k_output_chars=0.000,
+    ),
     # TODO(chengrun): Set a more appropriate rpm for endpoint.
     'vertexai-endpoint': pg.Dict(
         rpm=20,
@@ -215,7 +227,7 @@ class VertexAI(rest.REST):
     return (
         cost_per_1k_input_chars * num_input_tokens
         + cost_per_1k_output_chars * num_output_tokens
-    ) * AVGERAGE_CHARS_PER_TOEKN / 1000
+    ) * AVGERAGE_CHARS_PER_TOKEN / 1000
 
   @functools.cached_property
   def _session(self):
@@ -343,7 +355,7 @@ class VertexAI(rest.REST):
     return lf.AIMessage.from_chunks(chunks)
 
 
-_IMAGE_TYPES = [
+IMAGE_TYPES = [
     'image/png',
     'image/jpeg',
     'image/webp',
@@ -351,7 +363,7 @@ _IMAGE_TYPES = [
     'image/heif',
 ]
 
-_AUDIO_TYPES = [
+AUDIO_TYPES = [
     'audio/aac',
     'audio/flac',
     'audio/mp3',
@@ -362,10 +374,10 @@ _AUDIO_TYPES = [
     'audio/opus',
     'audio/pcm',
     'audio/wav',
-    'audio/webm'
+    'audio/webm',
 ]
 
-_VIDEO_TYPES = [
+VIDEO_TYPES = [
     'video/mov',
     'video/mpeg',
     'video/mpegps',
@@ -375,9 +387,10 @@ _VIDEO_TYPES = [
     'video/wmv',
     'video/x-flv',
     'video/3gpp',
+    'video/quicktime',
 ]
 
-_DOCUMENT_TYPES = [
+DOCUMENT_TYPES = [
     'application/pdf',
     'text/plain',
     'text/csv',
@@ -388,11 +401,25 @@ _DOCUMENT_TYPES = [
 ]
 
 
+class VertexAIGemini2_0(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 2.0 model."""
+
+  supported_modalities: pg.typing.List(str).freeze(  # pytype: disable=invalid-annotation
+      DOCUMENT_TYPES + IMAGE_TYPES + AUDIO_TYPES + VIDEO_TYPES
+  )
+
+
+class VertexAIGeminiFlash2_0Exp(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 2.0 Flash model."""
+
+  model = 'gemini-2.0-flash-exp'
+
+
 class VertexAIGemini1_5(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 model."""
 
-  supported_modalities: pg.typing.List(str).freeze(   # pytype: disable=invalid-annotation
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
+  supported_modalities: pg.typing.List(str).freeze(  # pytype: disable=invalid-annotation
+      DOCUMENT_TYPES + IMAGE_TYPES + AUDIO_TYPES + VIDEO_TYPES
   )
 
 
@@ -460,8 +487,8 @@ class VertexAIGeminiPro1Vision(VertexAI):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.0 Pro Vision model."""
 
   model = 'gemini-1.0-pro-vision'
-  supported_modalities: pg.typing.List(str).freeze(   # pytype: disable=invalid-annotation
-      _IMAGE_TYPES + _VIDEO_TYPES
+  supported_modalities: pg.typing.List(str).freeze(  # pytype: disable=invalid-annotation
+      IMAGE_TYPES + VIDEO_TYPES
   )
 
 

langfun/core/message.py

@@ -769,7 +769,6 @@ class Message(
             padding: 20px;
             margin: 10px 5px 10px 5px;
             font-style: italic;
-            font-size: 1.1em;
             white-space: pre-wrap;
             border: 1px solid #EEE;
             border-radius: 5px;
@@ -778,7 +777,7 @@ class Message(
         .modality-in-text {
             display: inline-block;
         }
-        .modality-in-text > details {
+        .modality-in-text > details.pyglove {
             display: inline-block;
             font-size: 0.8em;
             border: 0;

langfun/core/message_test.py

@@ -380,7 +380,6 @@ class MessageTest(unittest.TestCase):
                 padding: 20px;
                 margin: 10px 5px 10px 5px;
                 font-style: italic;
-                font-size: 1.1em;
                 white-space: pre-wrap;
                 border: 1px solid #EEE;
                 border-radius: 5px;
@@ -389,7 +388,7 @@ class MessageTest(unittest.TestCase):
             .modality-in-text {
                 display: inline-block;
             }
-            .modality-in-text > details {
+            .modality-in-text > details.pyglove {
                 display: inline-block;
                 font-size: 0.8em;
                 border: 0;