diff --git a/README.md b/README.md index 09499facc..1403e13d5 100644 --- a/README.md +++ b/README.md @@ -142,10 +142,15 @@ response.reply("Can you tell me a joke?") ### Colab magics -Once installed, use the Python client via the `%%palm` Colab magic. Read the [full guide](https://github.com/google/generative-ai-docs/blob/main/site/en/palm_docs/notebook_magic.ipynb). - ``` -%%palm +%pip install -q google-generativeai +%load_ext google.generativeai.notebook +``` + +Once installed, use the Python client via the `%%llm` Colab magic. Read the full guide [here](https://developers.generativeai.google/tools/notebook_magic). + +```python +%%llm The best thing since sliced bread is ``` diff --git a/google/generativeai/notebook/cmd_line_parser.py b/google/generativeai/notebook/cmd_line_parser.py index 854b5f5c7..7005bf600 100644 --- a/google/generativeai/notebook/cmd_line_parser.py +++ b/google/generativeai/notebook/cmd_line_parser.py @@ -368,7 +368,7 @@ def _create_parser( placeholders: AbstractSet[str] | None, ) -> argparse.ArgumentParser: """Create the full parser.""" - system_name = "palm" + system_name = "llm" description = "A system for interacting with LLMs." epilog = "" diff --git a/google/generativeai/notebook/magics.py b/google/generativeai/notebook/magics.py index 6fb5e22fa..8b6af51b9 100644 --- a/google/generativeai/notebook/magics.py +++ b/google/generativeai/notebook/magics.py @@ -14,14 +14,14 @@ # limitations under the License. """Colab Magics class. -Installs %%palm magics. +Installs %%llm magics. """ from __future__ import annotations import abc from google.auth import credentials -from google.generativeai import client as palm +from google.generativeai import client as genai from google.generativeai.notebook import gspread_client from google.generativeai.notebook import ipython_env from google.generativeai.notebook import ipython_env_impl @@ -34,9 +34,9 @@ # Set the UA to distinguish the magic from the client. Do this at import-time -# so that a user can still call `palm.configure()`, and both their settings +# so that a user can still call `genai.configure()`, and both their settings # and this are honored. -palm.USER_AGENT = "genai-py-magic" +genai.USER_AGENT = "genai-py-magic" SheetsInputs = sheets_utils.SheetsInputs SheetsOutputs = sheets_utils.SheetsOutputs @@ -72,7 +72,7 @@ class AbstractMagics(abc.ABC): """Defines interface to Magics class.""" @abc.abstractmethod - def palm(self, cell_line: str | None, cell_body: str | None): + def llm(self, cell_line: str | None, cell_body: str | None): """Perform various LLM-related operations. Args: @@ -92,7 +92,7 @@ class MagicsImpl(AbstractMagics): def __init__(self): self._engine = magics_engine.MagicsEngine(env=_get_ipython_env()) - def palm(self, cell_line: str | None, cell_body: str | None): + def llm(self, cell_line: str | None, cell_body: str | None): """Perform various LLM-related operations. Args: @@ -126,7 +126,7 @@ def get_instance(cls) -> AbstractMagics: return cls._instance @magic.line_cell_magic - def palm(self, cell_line: str | None, cell_body: str | None): + def llm(self, cell_line: str | None, cell_body: str | None): """Perform various LLM-related operations. Args: @@ -136,7 +136,15 @@ def palm(self, cell_line: str | None, cell_body: str | None): Returns: Results from running MagicsEngine. """ - return Magics.get_instance().palm(cell_line=cell_line, cell_body=cell_body) + return Magics.get_instance().llm(cell_line=cell_line, cell_body=cell_body) + + @magic.line_cell_magic + def palm(self, cell_line: str | None, cell_body: str | None): + return self.llm(cell_line, cell_body) + + @magic.line_cell_magic + def gemini(self, cell_line: str | None, cell_body: str | None): + return self.llm(cell_line, cell_body) IPython.get_ipython().register_magics(Magics) diff --git a/google/generativeai/notebook/text_model.py b/google/generativeai/notebook/text_model.py index dbe79097b..1d13a28aa 100644 --- a/google/generativeai/notebook/text_model.py +++ b/google/generativeai/notebook/text_model.py @@ -16,13 +16,15 @@ from __future__ import annotations from google.api_core import retry -from google.generativeai import text -from google.generativeai.types import text_types +import google.generativeai as genai +from google.generativeai.types import generation_types from google.generativeai.notebook.lib import model as model_lib +_DEFAULT_MODEL = "models/gemini-pro" + class TextModel(model_lib.AbstractModel): - """Concrete model that uses the Text service.""" + """Concrete model that uses the generate_content service.""" def _generate_text( self, @@ -30,15 +32,16 @@ def _generate_text( model: str | None = None, temperature: float | None = None, candidate_count: int | None = None, - **kwargs, - ) -> text_types.Completion: - if model is not None: - kwargs["model"] = model + ) -> generation_types.GenerateContentResponse: + gen_config = {} if temperature is not None: - kwargs["temperature"] = temperature + gen_config["temperature"] = temperature if candidate_count is not None: - kwargs["candidate_count"] = candidate_count - return text.generate_text(prompt=prompt, **kwargs) + gen_config["candidate_count"] = candidate_count + + model_name = model or _DEFAULT_MODEL + gen_model = genai.GenerativeModel(model_name=model_name) + return gen_model.generate_content(prompt, generation_config=gen_config) def call_model( self, @@ -58,7 +61,11 @@ def call_model( candidate_count=model_args.candidate_count, ) + text_outputs = [] + for c in response.candidates: + text_outputs.append("".join(p.text for p in c.content.parts)) + return model_lib.ModelResults( model_input=model_input, - text_results=[x["output"] for x in response.candidates], + text_results=text_outputs, ) diff --git a/tests/notebook/test_magics_engine.py b/tests/notebook/test_magics_engine.py index 493751d3e..726498b0d 100644 --- a/tests/notebook/test_magics_engine.py +++ b/tests/notebook/test_magics_engine.py @@ -248,7 +248,7 @@ def test_run_help(self): # Should not raise an exception. results = engine.execute_cell(magic_line, "ignored") - self.assertRegex(str(results), "usage: palm run") + self.assertRegex(str(results), "usage: llm run") def test_error(self): mock_registry = EchoModelRegistry() diff --git a/tests/notebook/text_model_test.py b/tests/notebook/text_model_test.py index d7cbe0924..9239ac9c3 100644 --- a/tests/notebook/text_model_test.py +++ b/tests/notebook/text_model_test.py @@ -19,7 +19,7 @@ from absl.testing import absltest from google.api_core import exceptions -from google.generativeai import text +from google.generativeai.types import generation_types from google.generativeai.notebook import text_model from google.generativeai.notebook.lib import model as model_lib @@ -29,20 +29,23 @@ def _fake_generator( model: str | None = None, temperature: float | None = None, candidate_count: int | None = None, -) -> text.Completion: - return text.Completion( - prompt=prompt, - model=model, - temperature=temperature, - candidate_count=candidate_count, - # Smuggle the parameters as text output, so we can make assertions. - candidates=[ - {"output": f"{prompt}_1"}, - {"output": model}, - {"output": temperature}, - {"output": candidate_count}, - ], - ) +): + def make_candidate(txt): + c = mock.Mock() + p = mock.Mock() + p.text = str(txt) + c.content.parts = [p] + return c + + response = mock.Mock() + # Smuggle the parameters as text output, so we can make assertions. + response.candidates = [ + make_candidate(f"{prompt}_1"), + make_candidate(model), + make_candidate(temperature), + make_candidate(candidate_count), + ] + return response class TestModel(text_model.TextModel): @@ -55,7 +58,7 @@ def _generate_text( temperature: float | None = None, candidate_count: int | None = None, **kwargs, - ) -> text.Completion: + ) -> generation_types.GenerateContentResponse: return _fake_generator( prompt=prompt, model=model,