From 8ed14b0ef6e825bad46a1098652ce73ae42c714a Mon Sep 17 00:00:00 2001
From: Meor Amer <92068895+mrmer1@users.noreply.github.com>
Date: Thu, 24 Oct 2024 23:09:02 +0800
Subject: [PATCH] update integration code snippets (#211)

Co-authored-by: trentfowlercohere <141260477+trentfowlercohere@users.noreply.github.com>
---
 .../chat-on-langchain.mdx                     |  4 +-
 fern/pages/integrations/llamaindex.mdx        | 89 ++++++++++++++-----
 2 files changed, 67 insertions(+), 26 deletions(-)

diff --git a/fern/pages/integrations/cohere-and-langchain/chat-on-langchain.mdx b/fern/pages/integrations/cohere-and-langchain/chat-on-langchain.mdx
index 2b8242e39..cadd73c15 100644
--- a/fern/pages/integrations/cohere-and-langchain/chat-on-langchain.mdx
+++ b/fern/pages/integrations/cohere-and-langchain/chat-on-langchain.mdx
@@ -30,14 +30,14 @@ llm = ChatCohere(cohere_api_key="COHERE_API_KEY",
 
 # Send a chat message without chat history
 current_message = [HumanMessage(content="knock knock")]
-print(llm(current_message))
+print(llm.invoke(current_message))
 
 # Send a chat message with chat history, note the last message is the current user message
 current_message_and_history = [
     HumanMessage(content="knock knock"),
     AIMessage(content="Who's there?"),
     HumanMessage(content="Tank") ]
-print(llm(current_message_and_history))
+print(llm.invoke(current_message_and_history))
 ```
 
 ### Cohere Agents with LangChain
diff --git a/fern/pages/integrations/llamaindex.mdx b/fern/pages/integrations/llamaindex.mdx
index b539b596d..100bcabb2 100644
--- a/fern/pages/integrations/llamaindex.mdx
+++ b/fern/pages/integrations/llamaindex.mdx
@@ -14,7 +14,11 @@ updatedAt: 'Tue May 14 2024 15:10:49 GMT+0000 (Coordinated Universal Time)'
 
 To use LlamaIndex and Cohere, you will need: 
 
-- LlamaIndex Package. To install it, run `pip install llama-index`, `llama-index-llms-cohere`, and `llama-index-embeddings-cohere`. 
+- LlamaIndex Package. To install it, run:
+    - `pip install llama-index`
+    - `pip install llama-index-llms-cohere` (to use the Command models)
+    - `pip install llama-index-embeddings-cohere` (to use the Embed models)
+    - `pip install llama-index-postprocessor-cohere-rerank` (to use the Rerank models)
 - Cohere's SDK. To install it, run `pip install cohere`. If you run into any issues or want more details on Cohere's SDK, [see this wiki](https://github.com/cohere-ai/cohere-python).
 - A Cohere API Key. For more details on pricing [see this page](https://cohere.com/pricing). When you create an account with Cohere, we automatically create a trial API key for you. This key will be available on the dashboard where you can copy it, and it's in the dashboard section called "API Keys" as well.
 
@@ -26,10 +30,13 @@ To use Cohere's chat functionality with LlamaIndex create a [Cohere model object
 from llama_index.llms.cohere import Cohere
 from llama_index.core.llms import ChatMessage
 
-cohere_model = Cohere(api_key=os.getenv("COHERE_API_KEY"))
+cohere_model = Cohere(api_key="COHERE_API_KEY",
+                      model="command-r-plus")
+
 message = ChatMessage(role="user", content= "What is 2 + 3?")
-resp = cohere_model.chat([message])
-print(resp)
+
+response = cohere_model.chat([message])
+print(response)
 ```
 
 ### Cohere Embeddings with LlamaIndex
@@ -40,9 +47,9 @@ To use Cohere's embeddings with LlamaIndex create a [Cohere Embeddings object](h
 from llama_index.embeddings.cohere import CohereEmbedding
 
 embed_model = CohereEmbedding(
-    cohere_api_key=cohere_api_key,
-    model_name="embed-english-v3.0", # Supports all Cohere embed models
-    input_type="search_query", # Required for v3 models
+    api_key="COHERE_API_KEY",
+    model_name="embed-english-v3.0",
+    input_type="search_document", # Use search_query for queries, search_document for documents
 )
 
 # Generate Embeddings
@@ -58,7 +65,38 @@ print(embeddings[:5])
 To use Cohere's rerank functionality with LlamaIndex create a [ Cohere Rerank object ](https://docs.llamaindex.ai/en/latest/examples/node_postprocessor/CohereRerank.html#) and use as a [node post processor.](https://docs.llamaindex.ai/en/stable/module_guides/querying/node_postprocessors/root.html)
 
 ```python PYTHON
-cohere_rerank = CohereRerank(api_key="{API_KEY}", top_n=2)
+from llama_index.postprocessor.cohere_rerank import CohereRerank
+from llama_index.readers.web import SimpleWebPageReader # first, run `pip install llama-index-readers-web`
+
+# create index (we are using an example page from Cohere's docs)
+documents = SimpleWebPageReader(html_to_text=True).load_data(
+    ["https://docs.cohere.com/v2/docs/prompt-tuner"]
+) # you can replace this with any other reader or documents
+index = VectorStoreIndex.from_documents(documents=documents)
+
+# create reranker
+cohere_rerank = CohereRerank(api_key="COHERE_API_KEY",
+                             model="rerank-english-v3.0",
+                             top_n=2)
+
+# query the index
+query_engine = index.as_query_engine(
+    similarity_top_k=10,
+    node_postprocessors=[cohere_rerank],
+)
+
+print(query_engine)
+
+# generate a response
+response = query_engine.query(
+    "What is Cohere Prompt Tuner?",
+)
+
+print(response)
+
+# To view the source documents
+from llama_index.core.response.pprint_utils import pprint_response
+pprint_response(response, show_source=True)
 ```
 
 ### Cohere RAG with LlamaIndex
@@ -66,39 +104,42 @@ cohere_rerank = CohereRerank(api_key="{API_KEY}", top_n=2)
 The following example uses Cohere's chat model, embeddings and rerank functionality to generate a response based on your data. 
 
 ```python PYTHON
-# rerank
-
-from llama_index import ServiceContext, VectorStoreIndex
 from llama_index.llms.cohere import Cohere
-from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
-from llama_index.embeddings.cohereai import CohereEmbedding
+from llama_index.embeddings.cohere import CohereEmbedding
 from llama_index.postprocessor.cohere_rerank import CohereRerank
+from llama_index.core import Settings
+from llama_index.core import VectorStoreIndex
+from llama_index.readers.web import SimpleWebPageReader # first, run `pip install llama-index-readers-web`
 
 # Create the embedding model
 embed_model = CohereEmbedding(
-    cohere_api_key="{API_KEY}",
+    api_key="COHERE_API_KEY",
     model_name="embed-english-v3.0",
     input_type="search_query",
 )
 
 # Create the service context with the cohere model for generation and embedding model
-service_context = ServiceContext.from_defaults(
-    llm=Cohere(api_key="{API_KEY}", model="command"),
-    embed_model=embed_model
-)
+Settings.llm = Cohere(api_key="COHERE_API_KEY",
+                      model="command-r-plus")
+Settings.embed_model = embed_model
 
-# Load the data, for this example data needs to be in a test file 
-data = SimpleDirectoryReader(input_files=["example_data_file.txt"]).load_data()
-index = VectorStoreIndex.from_documents(data, service_context=service_context)
+# create index (we are using an example page from Cohere's docs)
+documents = SimpleWebPageReader(html_to_text=True).load_data(
+    ["https://docs.cohere.com/v2/docs/prompt-tuner"]
+) # you can replace this with any other reader or documents
+index = VectorStoreIndex.from_documents(documents=documents)
 
 # Create a cohere reranker 
-cohere_rerank = CohereRerank(api_key="{API_KEY}")
+cohere_rerank = CohereRerank(api_key="COHERE_API_KEY",
+                             model="rerank-english-v3.0",
+                             top_n=2)
 
 # Create the query engine
 query_engine = index.as_query_engine(node_postprocessors=[cohere_rerank])
 
 # Generate the response
-response = query_engine.query("Who founded Cohere?",)
+response = query_engine.query("What is Cohere Prompt Tuner?")
+print(response)
 ```
 
 ### Cohere Tool Use (Function Calling) with LlamaIndex
@@ -124,7 +165,7 @@ def add(a: int, b: int) -> int:
 add_tool = FunctionTool.from_defaults(fn=add)
 
 # Define LLM
-llm = Cohere(api_key=os.getenv("COHERE_API_KEY"),
+llm = Cohere(api_key="COHERE_API_KEY",
              model="command-r-plus")
 
 # Create agent