Update agent.py

agent.py

@@ -1,226 +1,83 @@
-from langchain_huggingface import HuggingFacePipeline
 from transformers import pipeline
 import requests
 from bs4 import BeautifulSoup
-import os
-from typing import List, TypedDict
-from dataclasses import dataclass
-from langgraph.graph import StateGraph, START, END
-from dotenv import load_dotenv
 from time import sleep
 
-load_dotenv()
-
-@dataclass
-class Command:
-    update: dict = None
-    goto: str = None
-
-# Initialize local Hugging Face model with error handling
+# Load the model once
 try:
-    hf_pipeline = pipeline(
-        "text2text-generation",
-        model="google/flan-t5-small",
-        max_length=512,
-        temperature=0.7,
-    )
-    model = HuggingFacePipeline(pipeline=hf_pipeline)
+    hf_pipe = pipeline("text2text-generation", model="google/flan-t5-small", max_length=512, temperature=0.7)
 except Exception as e:
-    print(f"Error initializing model pipeline: {e}")
-    model = None  # fallback
-
-def scrape_startpage(query: str, max_results: int = 3) -> List[dict]:
-    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"
-    }
-    for attempt in range(3):
-        try:
-            response = requests.get(url, headers=headers, timeout=10)
-            response.raise_for_status()
-            soup = BeautifulSoup(response.text, "html.parser")
-            results = []
-            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")
-                title_text = title.get_text().strip() if title else "No title"
-                snippet_text = snippet.get_text().strip() if snippet else "No snippet"
-                results.append({"title": title_text, "snippet": snippet_text})
-            return results
-        except Exception as e:
-            print(f"Error scraping Startpage (attempt {attempt+1}/3): {e}")
-            sleep(2 ** attempt)
-    return []
+    print(f"Error loading model: {e}")
+    hf_pipe = None
 
 def safe_invoke(prompt: str) -> str:
-    if not model:
-        return "Error: model not initialized"
+    if not hf_pipe:
+        return "Error: Model not loaded"
     try:
-        response = model.invoke(prompt)
-        if not response or not isinstance(response, str):
-            return "Error: empty or invalid response"
-        return response.strip()
-    except Exception as e:
-        print(f"Error invoking model: {e}")
+        outputs = hf_pipe(prompt)
+        if outputs and isinstance(outputs, list):
+            return outputs[0]['generated_text'].strip()
         return "Error"
-
-def get_platform_tips(state) -> Command:
-    query = f"tips on how to write an effective post on {state.get('platform', '')}"
-    results = scrape_startpage(query)
-    if results:
-        prompt = f"Summarize these tips for writing effective posts on {state.get('platform', '')}: {results}"
-        response = safe_invoke(prompt)
-    else:
-        response = f"Use a professional tone, include a clear call-to-action, and keep the post concise for {state.get('platform', '')}."
-    return Command(update={"tips": response}, goto="web_search")
-
-def web_search(state) -> Command:
-    search_results = scrape_startpage(state.get("topic", ""), max_results=3)
-    return Command(update={"search_results": search_results}, goto="generate_post")
-
-def generate_social_media_post(state) -> Command:
-    try:
-        prompt = f"""
-        You are a social media strategist for a B2B bank. Generate a {state.get("platform", "")} post.
-        The post should:
-        - Be engaging but professional.
-        - Provide value to corporate clients.
-        - Focus on {state.get("topic", "")}.
-        - Incorporate information from {state.get("search_results", "general knowledge about the topic")}
-        Output as plain text.
-        """
-        response = safe_invoke(prompt)
-        return Command(update={"post": response}, goto="evaluate_engagement")
-    except Exception as e:
-        print(f"Error in generate_social_media_post: {e}")
-        return Command(update={"post": "Error generating post"}, goto="evaluate_engagement")
-
-def evaluate_engagement(state) -> Command:
-    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.get("platform", "")}
-    Post: {state.get("post", "")}
-    Respond with only a number between 1 and 10, no text.
-    """
-    score = safe_invoke(prompt)
-    return Command(update={"engagement_score": score}, goto="evaluate_tone")
-
-def evaluate_tone(state) -> Command:
-    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.get("platform", "")}
-    Post: {state.get("post", "")}
-    Respond with only a number between 1 and 10, no text.
-    """
-    score = safe_invoke(prompt)
-    return Command(update={"tone_score": score}, goto="evaluate_clarity")
-
-def evaluate_clarity(state) -> Command:
-    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.get("platform", "")}
-    Post: {state.get("post", "")}
-    Respond with only a number between 1 and 10, no text.
-    """
-    score = safe_invoke(prompt)
-    return Command(update={"clarity_score": score}, goto="revise_if_needed")
-
-def revise_if_needed(state) -> Command:
-    try:
-        scores = [int(state.get("engagement_score", "0")),
-                  int(state.get("tone_score", "0")),
-                  int(state.get("clarity_score", "0"))]
     except Exception as e:
-        print(f"Error parsing scores: {e}")
-        return Command(update={"post": "Error: Invalid scores"}, 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.get("post", "")}
-
-        Improve based on the following scores:
-        Engagement: {state.get("engagement_score", "0")}
-        Tone: {state.get("tone_score", "0")}
-        Clarity: {state.get("clarity_score", "0")}
-        """
-        revised_post = safe_invoke(prompt)
-        return Command(update={"post": revised_post}, goto="get_image")
-
-    return Command(goto="get_image")
-
-def fetch_image(state) -> Command:
-    pexels_key = os.getenv("PEXELS_API_KEY")
-    if not pexels_key:
-        print("PEXELS_API_KEY not set in environment")
-        return Command(update={"image_url": []}, goto=END)
-
-    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.
-    Topic: {state.get('topic', '')}
-    """
-    query = safe_invoke(prompt)
-    if "Error" in query or not query:
-        query = state.get('topic', '')
+        print(f"Error during generation: {e}")
+        return "Error"
 
-    url = "https://api.pexels.com/v1/search"
-    params = {
-        "query": query,
-        "per_page": 5,
-        "page": 1
-    }
-    headers = {
-        "Authorization": pexels_key
-    }
+def scrape_startpage(query: str, max_results: int = 3):
+    url = f"https://www.startpage.com/sp/search?query={query.replace(' ', '+')}"
+    headers = {"User-Agent": "Mozilla/5.0"}
     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', [])]
-            return Command(update={"image_url": urls}, goto=END)
+            res = requests.get(url, headers=headers, timeout=10)
+            res.raise_for_status()
+            soup = BeautifulSoup(res.text, "html.parser")
+            results = []
+            # Find divs with class "result" (Startpage search results)
+            for r in soup.find_all("div", class_="result")[:max_results]:
+                title = r.find("h3")
+                desc = r.find("p", class_="desc")
+                title_text = title.get_text(strip=True) if title else "No title"
+                desc_text = desc.get_text(strip=True) if desc else "No description"
+                results.append(f"{title_text}: {desc_text}")
+            return results
         except Exception as e:
-            print(f"Error fetching images from Pexels (attempt {attempt+1}/3): {e}")
+            print(f"Scrape error (attempt {attempt+1}): {e}")
             sleep(2 ** attempt)
-    return Command(update={"image_url": []}, goto=END)
-
-class State(TypedDict, total=False):
-    topic: str
-    platform: str
-    tips: str
-    search_results: List[dict]
-    post: 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)
-workflow.add_node("evaluate_engagement", evaluate_engagement)
-workflow.add_node("evaluate_tone", evaluate_tone)
-workflow.add_node("evaluate_clarity", evaluate_clarity)
-workflow.add_node("revise_if_needed", revise_if_needed)
-workflow.add_node("get_image", fetch_image)
-
-workflow.add_edge(START, "get_tips")
-workflow.add_edge("get_tips", "web_search")
-workflow.add_edge("web_search", "generate_post")
-workflow.add_edge("generate_post", "evaluate_engagement")
-workflow.add_edge("evaluate_engagement", "evaluate_tone")
-workflow.add_edge("evaluate_tone", "evaluate_clarity")
-workflow.add_edge("evaluate_clarity", "revise_if_needed")
-workflow.add_edge("revise_if_needed", "get_image")
+    return []
 
-graph = workflow.compile()
+def generate_post(topic, platform, search_results):
+    base_prompt = f"""You are a social media expert. Write a professional {platform} post about "{topic}".
+Use this information to help you: {search_results}
+Make the post clear, engaging, and suitable for corporate clients.
+Output only the post text."""
+    return safe_invoke(base_prompt)
+
+def score_post(post, platform, score_type):
+    prompt = f"""Rate the following post on {score_type} from 1 to 10 (just give a number):
+Platform: {platform}
+Post: {post}
+"""
+    return safe_invoke(prompt)
+
+def workflow(topic, platform):
+    # Step 1: Web search results
+    search_results = scrape_startpage(topic)
+    combined_results = " | ".join(search_results) if search_results else "No additional info."
+
+    # Step 2: Generate post
+    post = generate_post(topic, platform, combined_results)
+    if post == "Error":
+        return post, "Error", "Error", "Error"
+
+    # Step 3: Scores
+    engagement = score_post(post, platform, "engagement")
+    tone = score_post(post, platform, "tone")
+    clarity = score_post(post, platform, "clarity")
+
+    # Validate scores (should be digits)
+    def valid_score(s):
+        return s and s.strip().isdigit()
+
+    if not all(map(valid_score, (engagement, tone, clarity))):
+        return post, "Error", "Error", "Error"
+
+    return post, engagement, tone, clarity