Polish style

langchain-ai · Sep 19, 2024 · 46a181b · 46a181b
2 parents 3477826 + 8e4412a
commit 46a181b
Show file tree

Hide file tree

Showing 11 changed files with 195 additions and 208 deletions.
diff --git a/.env.example b/.env.example
@@ -2,6 +2,8 @@ TAVILY_API_KEY=...
 
 # To separate your traces from other application
 LANGCHAIN_PROJECT=data-enrichment
+# LANGCHAIN_API_KEY=...
+# LANGCHAIN_TRACING_V2=true
 
 # The following depend on your selected configuration
 

diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml
@@ -50,4 +50,7 @@ jobs:
           path: src/
 
       - name: Run tests
+        env:
+          ANTHROPIC_API_KEY: afakekey
+          TAVILY_API_KEY: anotherfakekey
         run: yarn test
diff --git a/README.md b/README.md
@@ -93,50 +93,74 @@ End setup instructions
 
 3. Consider a research topic and desired extraction schema.
 
-As an example, here is a research topic we can consider.
+As an example, here is a research topic we can consider:
+
+```
+"Autonomous agents"
+```
+
+With an `extractionSchema` of:
+
+```json
+{
+  "type": "object",
+  "properties": {
+    "facts": {
+      "type": "array",
+      "description": "An array of facts retrieved from the provided sources",
+      "items": {
+        "type": "string"
+      }
+    }
+  },
+  "required": ["facts"]
+}
+```
+
+Another example topic with a more complex schema is:
 
 ```
 "Top 5 chip providers for LLM Training"
 ```
 
-And here is a desired extraction schema.
+And here is a desired `extractionSchema`:
 
 ```json
-"extractionSchema": {
-    "type": "object",
-    "properties": {
-        "companies": {
-            "type": "array",
-            "items": {
-                "type": "object",
-                "properties": {
-                    "name": {
-                        "type": "string",
-                        "description": "Company name"
-                    },
-                    "technologies": {
-                        "type": "string",
-                        "description": "Brief summary of key technologies used by the company"
-                    },
-                    "market_share": {
-                        "type": "string",
-                        "description": "Overview of market share for this company"
-                    },
-                    "future_outlook": {
-                        "type": "string",
-                        "description": "Brief summary of future prospects and developments in the field for this company"
-                    },
-                    "key_powers": {
-                        "type": "string",
-                        "description": "Which of the 7 Powers (Scale Economies, Network Economies, Counter Positioning, Switching Costs, Branding, Cornered Resource, Process Power) best describe this company's competitive advantage"
-                    }
-                },
-                "required": ["name", "technologies", "market_share", "future_outlook"]
-            },
-            "description": "List of companies"
-        }
-    },
-    "required": ["companies"]
+{
+  "type": "object",
+  "properties": {
+    "companies": {
+      "type": "array",
+      "items": {
+        "type": "object",
+        "properties": {
+          "name": {
+            "type": "string",
+            "description": "Company name"
+          },
+          "technologies": {
+            "type": "string",
+            "description": "Brief summary of key technologies used by the company"
+          },
+          "market_share": {
+            "type": "string",
+            "description": "Overview of market share for this company"
+          },
+          "future_outlook": {
+            "type": "string",
+            "description": "Brief summary of future prospects and developments in the field for this company"
+          },
+          "key_powers": {
+            "type": "string",
+            "description": "Which of the 7 Powers (Scale Economies, Network Economies, Counter Positioning, Switching Costs, Branding, Cornered Resource, Process Power) best describe this company's competitive advantage"
+          }
+        },
+        "required": ["name", "technologies", "market_share", "future_outlook"]
+      },
+      "description": "List of companies"
+    }
+  },
+  "required": ["companies"]
 }
 ```
 
@@ -145,7 +169,7 @@ And here is a desired extraction schema.
 ## How to customize
 
 1. **Customize research targets**: Provide a custom JSON `extractionSchema` when calling the graph to gather different types of information.
-2. **Select a different model**: We default to anthropic (claude-3-5-sonnet-20240620). You can select a compatible chat model using `provider/model-name` via configuration. Example: `openai/gpt-4o-mini`.
+2. **Select a different model**: We default to anthropic (`claude-3-5-sonnet-20240620`). You can select a compatible chat model using `provider/model-name` via configuration. Example: `openai/gpt-4o-mini`.
 3. **Customize the prompt**: We provide a default prompt in [src/enrichment_agent/prompts.ts](./src/enrichment_agent/prompts.ts). You can easily update this via configuration.
 
 For quick prototyping, these configurations can be set in the studio UI.
@@ -163,7 +187,7 @@ While iterating on your graph, you can edit past state and rerun your app from p
 
 Follow up requests will be appended to the same thread. You can create an entirely new thread, clearing previous history, using the `+` button in the top right.
 
-You can find the latest (under construction) docs on [LangGraph.JS](https://langchain-ai.github.io/langgraphjs/) here, including examples and other references. Using those guides can help you pick the right patterns to adapt here for your use case.
+You can find the latest (under construction) docs on [LangGraph.js](https://langchain-ai.github.io/langgraphjs/) here, including examples and other references. Using those guides can help you pick the right patterns to adapt here for your use case.
 
 LangGraph Studio also integrates with [LangSmith](https://smith.langchain.com/) for more in-depth tracing and collaboration with teammates.
 

diff --git a/package.json b/package.json
@@ -23,7 +23,7 @@
   "dependencies": {
     "@langchain/anthropic": "^0.3.1",
     "@langchain/community": "^0.3.1",
-    "@langchain/core": "^0.3.2",
+    "@langchain/core": "^0.3.3",
     "@langchain/langgraph": "^0.2.8",
     "langchain": "^0.3.2",
     "langsmith": "^0.1.59",
@@ -36,6 +36,7 @@
     "@jest/globals": "^29.7.0",
     "@tsconfig/recommended": "^1.0.7",
     "@types/jest": "^29.5.0",
+    "@types/node": "^20.14.8",
     "@typescript-eslint/eslint-plugin": "^5.59.8",
     "@typescript-eslint/parser": "^5.59.8",
     "dotenv": "^16.4.5",

diff --git a/src/enrichment_agent/configuration.ts b/src/enrichment_agent/configuration.ts
@@ -49,15 +49,15 @@ export const ConfigurationAnnotation = Annotation.Root({
 export function ensureConfiguration(
   config?: RunnableConfig,
 ): typeof ConfigurationAnnotation.State {
-  const configurable = (config?.configurable || {}) as Partial<
+  const configurable = (config?.configurable ?? {}) as Partial<
     typeof ConfigurationAnnotation.State
   >;
 
   return {
-    model: configurable.model || "anthropic/claude-3-5-sonnet-20240620",
-    prompt: configurable.prompt || MAIN_PROMPT,
-    maxSearchResults: configurable.maxSearchResults || 10,
-    maxInfoToolCalls: configurable.maxInfoToolCalls || 3,
-    maxLoops: configurable.maxLoops || 6,
+    model: configurable.model ?? "anthropic/claude-3-5-sonnet-20240620",
+    prompt: configurable.prompt ?? MAIN_PROMPT,
+    maxSearchResults: configurable.maxSearchResults ?? 5,
+    maxInfoToolCalls: configurable.maxInfoToolCalls ?? 3,
+    maxLoops: configurable.maxLoops ?? 6,
   };
 }
diff --git a/src/enrichment_agent/graph.ts b/src/enrichment_agent/graph.ts
@@ -14,12 +14,13 @@ import { RunnableConfig } from "@langchain/core/runnables";
 import { tool } from "@langchain/core/tools";
 import { StateGraph } from "@langchain/langgraph";
 import { z } from "zod";
+
 import {
   ConfigurationAnnotation,
   ensureConfiguration,
 } from "./configuration.js";
 import { AnyRecord, InputStateAnnotation, StateAnnotation } from "./state.js";
-import { toolNode, TOOLS } from "./tools.js";
+import { MODEL_TOOLS, toolNode } from "./tools.js";
 import { loadChatModel } from "./utils.js";
 
 /**
@@ -43,15 +44,13 @@ import { loadChatModel } from "./utils.js";
 async function callAgentModel(
   state: typeof StateAnnotation.State,
   config: RunnableConfig,
-): Promise<{
-  messages: BaseMessage[];
-  info?: AnyRecord;
-  loopStep: number;
-}> {
+): Promise<typeof StateAnnotation.Update> {
   const configuration = ensureConfiguration(config);
   // First, define the info tool. This uses the user-provided
   // json schema to define the research targets
-  const infoTool = tool(async (_args: AnyRecord) => {}, {
+  // We pass an empty function because we will not actually invoke this tool.
+  // We are just using it for formatting.
+  const infoTool = tool(async () => {}, {
     name: "Info",
     description: "Call this when you have gathered all the relevant info",
     schema: state.extractionSchema,
@@ -61,7 +60,7 @@ async function callAgentModel(
   if (!rawModel.bindTools) {
     throw new Error("Chat model does not support tool binding");
   }
-  const model = rawModel.bindTools([...TOOLS, infoTool], {
+  const model = rawModel.bindTools([...MODEL_TOOLS, infoTool], {
     tool_choice: "any",
   });
 
@@ -73,7 +72,7 @@ async function callAgentModel(
 
   // Next, we'll call the model.
   const response: AIMessage = await model.invoke(messages);
-  const response_messages = [response];
+  const responseMessages = [response];
 
   // If the model has collected enough information to fill uot
   // the provided schema, great! It will call the "Info" tool
@@ -96,13 +95,13 @@ async function callAgentModel(
     }
   } else {
     // If LLM didn't respect the tool_choice
-    response_messages.push(
+    responseMessages.push(
       new HumanMessage("Please respond by calling one of the provided tools."),
     );
   }
 
   return {
-    messages: response_messages,
+    messages: responseMessages,
     info,
     // This increments the step counter.
     // We configure a max step count to avoid infinite research loops
@@ -187,7 +186,7 @@ If you don't think it is good, you should be very specific about what could be i
   );
   messages.push({ role: "user", content: p1 });
 
-  // Calll the model
+  // Call the model
   const response = await boundModel.invoke(messages);
   if (response.is_satisfactory && presumedInfo) {
     return {
@@ -197,7 +196,7 @@ If you don't think it is good, you should be very specific about what could be i
           tool_call_id: lastMessage.tool_calls?.[0]?.id || "",
           content: response.reason.join("\n"),
           name: "Info",
-          additional_kwargs: { artifact: response },
+          artifact: response,
           status: "success",
         }),
       ],
@@ -209,7 +208,7 @@ If you don't think it is good, you should be very specific about what could be i
           tool_call_id: lastMessage.tool_calls?.[0]?.id || "",
           content: `Unsatisfactory response:\n${response.improvement_instructions}`,
           name: "Info",
-          additional_kwargs: { artifact: response },
+          artifact: response,
           status: "error",
         }),
       ],

diff --git a/src/enrichment_agent/state.ts b/src/enrichment_agent/state.ts
@@ -1,6 +1,5 @@
 import { Annotation, messagesStateReducer } from "@langchain/langgraph";
 import { type BaseMessage } from "@langchain/core/messages";
-import { z } from "zod";
 
 // eslint-disable-next-line
 export type AnyRecord = Record<string, any>;
@@ -11,11 +10,11 @@ export const InputStateAnnotation = Annotation.Root({
    * The info state trackes the current extracted data for the given topic,
    * conforming to the provided schema.
    */
-  info: Annotation<z.infer<z.ZodObject<z.ZodRawShape>>>,
+  info: Annotation<AnyRecord>,
   /**
    * The schema defines the information the agent is tasked with filling out.
    */
-  extractionSchema: Annotation<z.ZodObject<z.ZodRawShape>>,
+  extractionSchema: Annotation<AnyRecord>,
   // Feel free to add additional attributes to your state as needed.
   // Common examples include retrieved documents, extracted entities, API connections, etc.
 });
@@ -60,17 +59,17 @@ export const StateAnnotation = Annotation.Root({
     reducer: messagesStateReducer,
     default: () => [],
   }),
+
   topic: Annotation<string>,
   /**
    * The info state trackes the current extracted data for the given topic,
    * conforming to the provided schema.
    */
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
   info: Annotation<AnyRecord>,
+
   /**
    * The schema defines the information the agent is tasked with filling out.
    */
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
   extractionSchema: Annotation<AnyRecord>,
 
   /**