Update agent.py

agent.py

@@ -1,22 +1,23 @@
-from langchain_huggingface import HuggingFacePipeline
-from transformers import pipeline
-import requests
-from bs4 import BeautifulSoup
 import os
-from typing import List, TypedDict
+import requests
+from time import sleep
 from dataclasses import dataclass
-from langgraph.graph import StateGraph, START, END
+from typing import List, Optional, TypedDict
+
+from transformers import pipeline
+from langchain_huggingface import HuggingFacePipeline
+from bs4 import BeautifulSoup
 from dotenv import load_dotenv
-from time import sleep
+from langgraph.graph import StateGraph, START, END
 
 load_dotenv()
 
 @dataclass
 class Command:
-    update: dict = None
-    goto: str = None
+    update: Optional[dict] = None
+    goto: Optional[str] = None
 
-# Initialize local Hugging Face model
+# Initialize HF pipeline and wrap
 hf_pipeline = pipeline(
     "text2text-generation",
     model="google/flan-t5-small",
@@ -26,18 +27,19 @@ hf_pipeline = pipeline(
 model = HuggingFacePipeline(pipeline=hf_pipeline)
 
 def scrape_startpage(query: str, max_results: int = 3) -> List[dict]:
-    """Scrape search results from Startpage."""
     url = f"https://www.startpage.com/sp/search?query={query.replace(' ', '+')}"
     headers = {
-        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+        "User-Agent": (
+            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+            "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+        )
     }
     for attempt in range(3):
         try:
-            response = requests.get(url, headers=headers, timeout=10)
-            response.raise_for_status()
-            soup = BeautifulSoup(response.text, "html.parser")
+            r = requests.get(url, headers=headers, timeout=10)
+            r.raise_for_status()
+            soup = BeautifulSoup(r.text, "html.parser")
             results = []
-            # Extract search result snippets
             for result in soup.find_all("div", class_="result")[:max_results]:
                 title = result.find("h3") or result.find("a")
                 snippet = result.find("p", class_="desc")
@@ -46,159 +48,118 @@ def scrape_startpage(query: str, max_results: int = 3) -> List[dict]:
                 results.append({"title": title_text, "snippet": snippet_text})
             return results
         except Exception as e:
-            print(f"Error scraping Startpage (attempt {attempt+1}/3): {str(e)}")
+            print(f"Scrape error (try {attempt+1}): {e}")
             if attempt < 2:
-                sleep(2 ** attempt)  # Exponential backoff
+                sleep(2 ** attempt)
                 continue
             return []
 
 def get_platform_tips(state) -> Command:
-    """Scrape tips on writing effective posts for the provided platform from Startpage."""
     query = f"tips on how to write an effective post on {state['platform']}"
-    search_results = scrape_startpage(query, max_results=3)
-    if search_results:
-        prompt = f"""
-        Summarize the tips provided in {search_results}. These tips will be used to generate a {state['platform']} post.
-        Output as plain text.
-        """
+    results = scrape_startpage(query)
+    if results:
+        prompt = f"Summarize tips for {state['platform']} post: {results}. Output plain text."
         response = model.invoke(prompt)
     else:
-        response = f"Use a professional tone, include a clear call-to-action, and keep the post concise for {state['platform']}."
+        response = f"Write a concise professional post with a call-to-action on {state['platform']}."
     return Command(update={"tips": response}, goto="web_search")
 
 def web_search(state) -> Command:
-    """Scrape up-to-date information about the provided topic from Startpage."""
-    search_results = scrape_startpage(state["topic"], max_results=3)
-    return Command(update={"search_results": search_results}, goto="generate_post")
+    results = scrape_startpage(state["topic"])
+    return Command(update={"search_results": results}, goto="generate_post")
 
 def generate_social_media_post(state) -> Command:
-    """Generate a social media post for a B2B bank."""
     prompt = f"""
-    You are a social media strategist for a B2B bank. Generate a {state["platform"]} post.
-    The post should:
-    - Be engaging but professional.
-    - Provide value to corporate clients.
-    - Focus on {state["topic"]}.
-    - Incorporate information from {state["search_results"] if state["search_results"] else "general knowledge about the topic"}
-    Output as plain text.
+    Create a {state['platform']} post on {state['topic']} for B2B bank clients.
+    Use info: {state.get('search_results', [])}
+    Make it professional and engaging. Output text only.
     """
     response = model.invoke(prompt)
     return Command(update={"post": response}, goto="evaluate_engagement")
 
 def evaluate_engagement(state) -> Command:
-    """Assess how engaging the post is."""
     prompt = f"""
-    Score the following post on engagement (1-10) based on the provided social media platform.
-    Consider clarity, readability, and compelling call-to-action.
-    
-    Platform: {state["platform"]}
-    Post: {state["post"]}
-    
-    Respond with only a number between 1 and 10, no text.
+    Rate engagement 1-10 for post on {state['platform']}:
+    {state['post']}
+    Reply only a number.
     """
     score = model.invoke(prompt).strip()
     return Command(update={"engagement_score": score}, goto="evaluate_tone")
 
 def evaluate_tone(state) -> Command:
-    """Check if the post maintains a professional yet engaging tone."""
     prompt = f"""
-    Score the post’s tone (1-10). Ensure it's:
-    - Professional but not too rigid.
-    - Trustworthy and aligned with B2B financial services.
-    - Aligns with the specified platform.
-    Platform: {state["platform"]}
-    Post: {state["post"]}
-    
-    Respond with only a number between 1 and 10, no text.
+    Rate tone 1-10 (professional/trustworthy) for post:
+    {state['post']}
+    Reply only a number.
     """
     score = model.invoke(prompt).strip()
     return Command(update={"tone_score": score}, goto="evaluate_clarity")
 
 def evaluate_clarity(state) -> Command:
-    """Ensure the post is clear and not overly technical."""
     prompt = f"""
-    Score the post on clarity (1-10).
-    - Avoids jargon.
-    - Easy to read for busy corporate professionals.
-    - Appropriate for the social media platform.
-    Platform: {state["platform"]}
-    Post: {state["post"]}
-    
-    Respond with only a number between 1 and 10, no text.
+    Rate clarity 1-10 for post:
+    {state['post']}
+    Reply only a number.
     """
     score = model.invoke(prompt).strip()
     return Command(update={"clarity_score": score}, goto="revise_if_needed")
 
 def revise_if_needed(state) -> Command:
-    """Revise post if average evaluation score is below a threshold."""
     try:
-        scores = [int(state["engagement_score"]), int(state["tone_score"]), int(state["clarity_score"])]
+        scores = [int(state.get("engagement_score", "0")),
+                  int(state.get("tone_score", "0")),
+                  int(state.get("clarity_score", "0"))]
     except ValueError:
-        return Command(update={"post": "Error: Non-numeric scores received."}, goto="get_image")
-    avg_score = sum(scores) / len(scores)
-    
-    if avg_score < 7:
-        prompt = f"""
-        Revise this post to improve clarity, engagement, and tone:
-        
-        {state["post"]}
-        
-        Improve based on the following scores:
-        Engagement: {state["engagement_score"]}
-        Tone: {state["tone_score"]}
-        Clarity: {state["clarity_score"]}
-        """
-        revised_post = model.invoke(prompt)
-        return Command(update={"post": revised_post}, goto="get_image")
-    
-    return Command(goto="get_image")
+        return Command(goto="get_image")
+
+    avg = sum(scores) / 3
+    if avg < 7:
+        prompt = f"Improve post clarity, engagement, and tone:\n{state['post']}\nScores: {scores}"
+        revised = model.invoke(prompt)
+        return Command(update={"post": revised}, goto="get_image")
+    else:
+        return Command(goto="get_image")
 
 def fetch_image(state) -> Command:
-    """Fetch an image from Pexels based on the provided text."""
-    prompt = f"""
-    You are a search optimization assistant. Your task is to take a topic and improve it to ensure the best image results from an image search API like Pexels. Follow these steps:
-    1. Normalize the input: Convert all text to lowercase and remove special characters (except for spaces).
-    2. Add more descriptive terms: If the query is broad (e.g., "nature"), add more specific keywords like "landscape" or "outdoor".
-    3. Use synonyms and related terms: For terms with multiple meanings, add variations (e.g., "car" -> "vehicle automobile").
-    4. Specify style and tone: Add words like "professional", "modern", or "corporate" for B2B contexts.
-    5. Categorize the query: Add related terms for domains like "business", "finance", or "corporate".
-    Topic: {state['topic']}
-    """
+    api_key = os.getenv("PEXELS_API_KEY")
+    if not api_key:
+        print("Pexels API key missing.")
+        return Command(update={"image_url": []}, goto=END)
+
+    prompt = f"Generate a descriptive, professional image search query for: {state['topic']}"
+    search_query = model.invoke(prompt).strip()
+
     url = "https://api.pexels.com/v1/search"
-    params = {
-        "query": model.invoke(prompt).strip(),
-        "per_page": 5,
-        "page": 1
-    }
-    headers = {
-        "Authorization": os.getenv("PEXELS_API_KEY")
-    }
+    params = {"query": search_query, "per_page": 5}
+    headers = {"Authorization": api_key}
+
     for attempt in range(3):
         try:
-            response = requests.get(url, headers=headers, params=params)
-            response.raise_for_status()
-            data = response.json()
-            urls = [photo['url'] for photo in data.get('photos', [])]
+            r = requests.get(url, headers=headers, params=params)
+            r.raise_for_status()
+            data = r.json()
+            urls = [p["url"] for p in data.get("photos", [])]
             return Command(update={"image_url": urls}, goto=END)
-        except requests.RequestException as e:
-            print(f"Error fetching images from Pexels (attempt {attempt+1}/3): {e}")
+        except Exception as e:
+            print(f"Pexels request failed (try {attempt+1}): {e}")
             if attempt < 2:
-                sleep(2 ** attempt)  # Exponential backoff
+                sleep(2 ** attempt)
                 continue
             return Command(update={"image_url": []}, goto=END)
 
-class State(TypedDict):
+class State(TypedDict, total=False):
     topic: str
     platform: str
     tips: str
     search_results: List[dict]
     post: str
-    engagement_score: int
-    tone_score: int
-    clarity_score: int
-    image_url: str
+    engagement_score: str
+    tone_score: str
+    clarity_score: str
+    image_url: List[str]
 
 workflow = StateGraph(State)
+
 workflow.add_node("get_tips", get_platform_tips)
 workflow.add_node("web_search", web_search)
 workflow.add_node("generate_post", generate_social_media_post)
@@ -217,4 +178,4 @@ workflow.add_edge("evaluate_tone", "evaluate_clarity")
 workflow.add_edge("evaluate_clarity", "revise_if_needed")
 workflow.add_edge("revise_if_needed", "get_image")
 
-graph = workflow.compile()
+graph = workflow.compile()